airkjw/CLIProxyAPI
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add request logging capabilities to API handlers and update .gitignore
docker-image / docker (push) Has been cancelled
Details
goreleaser / goreleaser (push) Has been cancelled
Details

Browse Source 
Enhance API response handling by storing responses in context and updating request logger to include API responses
This commit is contained in:
Luis Pater
2025-08-16 05:03:10 +08:00
parent 4155805ad6
commit fcadf08921
10 changed files with 191 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/api/handlers/gemini/cli/cli_handlers.go
+34 -6
View File
@@ -127,6 +127,7 @@ func (h *GeminiCLIAPIHandlers) CLIHandler(c *gin.Context) {
return
}
_, _ = c.Writer.Write(output)
c.Set("API_RESPONSE", output)
}
}
@@ -155,7 +156,9 @@ func (h *GeminiCLIAPIHandlers) handleInternalStreamGenerateContent(c *gin.Contex
modelResult := gjson.GetBytes(rawJSON, "model")
modelName := modelResult.String()
cliCtx, cliCancel := context.WithCancel(context.Background())
backgroundCtx, cliCancel := context.WithCancel(context.Background())
cliCtx := context.WithValue(backgroundCtx, "gin", c)
var cliClient client.Client
defer func() {
// Ensure the client's mutex is unlocked on function exit.
@@ -184,21 +187,28 @@ outLoop:
// Send the message and receive response chunks and errors via channels.
respChan, errChan := cliClient.SendRawMessageStream(cliCtx, rawJSON, "")
hasFirstResponse := false
apiResponseData := make([]byte, 0)
for {
select {
// Handle client disconnection.
case <-c.Request.Context().Done():
if c.Request.Context().Err().Error() == "context canceled" {
log.Debugf("GeminiClient disconnected: %v", c.Request.Context().Err())
c.Set("API_RESPONSE", apiResponseData)
cliCancel() // Cancel the backend request.
return
}
// Process incoming response chunks.
case chunk, okStream := <-respChan:
if !okStream {
c.Set("API_RESPONSE", apiResponseData)
cliCancel()
return
}
apiResponseData = append(apiResponseData, chunk...)
hasFirstResponse = true
if cliClient.(*client.GeminiClient).GetGenerativeLanguageAPIKey() != "" {
chunk, _ = sjson.SetRawBytes(chunk, "response", chunk)
@@ -217,6 +227,7 @@ outLoop:
c.Status(err.StatusCode)
_, _ = fmt.Fprint(c.Writer, err.Error.Error())
flusher.Flush()
c.Set("API_RESPONSE", []byte(err.Error.Error()))
cliCancel()
}
return
@@ -237,7 +248,9 @@ func (h *GeminiCLIAPIHandlers) handleInternalGenerateContent(c *gin.Context, raw
// log.Debugf("GenerateContent: %s", string(rawJSON))
modelResult := gjson.GetBytes(rawJSON, "model")
modelName := modelResult.String()
cliCtx, cliCancel := context.WithCancel(context.Background())
backgroundCtx, cliCancel := context.WithCancel(context.Background())
cliCtx := context.WithValue(backgroundCtx, "gin", c)
var cliClient client.Client
defer func() {
if cliClient != nil {
@@ -269,11 +282,13 @@ func (h *GeminiCLIAPIHandlers) handleInternalGenerateContent(c *gin.Context, raw
c.Status(err.StatusCode)
_, _ = c.Writer.Write([]byte(err.Error.Error()))
log.Debugf("code: %d, error: %s", err.StatusCode, err.Error.Error())
c.Set("API_RESPONSE", []byte(err.Error.Error()))
cliCancel()
}
break
} else {
_, _ = c.Writer.Write(resp)
c.Set("API_RESPONSE", resp)
cliCancel()
break
}
@@ -313,7 +328,9 @@ func (h *GeminiCLIAPIHandlers) handleCodexInternalStreamGenerateContent(c *gin.C
modelName := gjson.GetBytes(rawJSON, "model")
cliCtx, cliCancel := context.WithCancel(context.Background())
backgroundCtx, cliCancel := context.WithCancel(context.Background())
cliCtx := context.WithValue(backgroundCtx, "gin", c)
var cliClient client.Client
defer func() {
// Ensure the client's mutex is unlocked on function exit.
@@ -345,24 +362,28 @@ outLoop:
ResponseID: "",
LastStorageOutput: "",
}
apiResponseData := make([]byte, 0)
for {
select {
// Handle client disconnection.
case <-c.Request.Context().Done():
if c.Request.Context().Err().Error() == "context canceled" {
log.Debugf("CodexClient disconnected: %v", c.Request.Context().Err())
c.Set("API_RESPONSE", apiResponseData)
cliCancel() // Cancel the backend request.
return
}
// Process incoming response chunks.
case chunk, okStream := <-respChan:
if !okStream {
c.Set("API_RESPONSE", apiResponseData)
cliCancel()
return
}
// _, _ = logFile.Write(chunk)
// _, _ = logFile.Write([]byte("\n"))
apiResponseData = append(apiResponseData, chunk...)
if bytes.HasPrefix(chunk, []byte("data: ")) {
jsonData := chunk[6:]
data := gjson.ParseBytes(jsonData)
@@ -390,6 +411,7 @@ outLoop:
c.Status(errMessage.StatusCode)
_, _ = fmt.Fprint(c.Writer, errMessage.Error.Error())
flusher.Flush()
c.Set("API_RESPONSE", []byte(errMessage.Error.Error()))
cliCancel()
}
return
@@ -416,7 +438,9 @@ func (h *GeminiCLIAPIHandlers) handleCodexInternalGenerateContent(c *gin.Context
modelName := gjson.GetBytes(rawJSON, "model")
cliCtx, cliCancel := context.WithCancel(context.Background())
backgroundCtx, cliCancel := context.WithCancel(context.Background())
cliCtx := context.WithValue(backgroundCtx, "gin", c)
var cliClient client.Client
defer func() {
// Ensure the client's mutex is unlocked on function exit.
@@ -440,22 +464,25 @@ outLoop:
// Send the message and receive response chunks and errors via channels.
respChan, errChan := cliClient.SendRawMessageStream(cliCtx, []byte(newRequestJSON), "")
apiResponseData := make([]byte, 0)
for {
select {
// Handle client disconnection.
case <-c.Request.Context().Done():
if c.Request.Context().Err().Error() == "context canceled" {
log.Debugf("CodexClient disconnected: %v", c.Request.Context().Err())
c.Set("API_RESPONSE", apiResponseData)
cliCancel() // Cancel the backend request.
return
}
// Process incoming response chunks.
case chunk, okStream := <-respChan:
if !okStream {
c.Set("API_RESPONSE", apiResponseData)
cliCancel()
return
}
apiResponseData = append(apiResponseData, chunk...)
if bytes.HasPrefix(chunk, []byte("data: ")) {
jsonData := chunk[6:]
data := gjson.ParseBytes(jsonData)
@@ -479,6 +506,7 @@ outLoop:
log.Debugf("org: %s", string(orgRawJSON))
log.Debugf("raw: %s", string(rawJSON))
log.Debugf("newRequestJSON: %s", newRequestJSON)
c.Set("API_RESPONSE", []byte(err.Error.Error()))
cliCancel()
}
return
internal/api/handlers/gemini/gemini_handlers.go
+37 -5
View File
@@ -267,7 +267,9 @@ func (h *GeminiAPIHandlers) handleGeminiStreamGenerateContent(c *gin.Context, ra
modelResult := gjson.GetBytes(rawJSON, "model")
modelName := modelResult.String()
cliCtx, cliCancel := context.WithCancel(context.Background())
backgroundCtx, cliCancel := context.WithCancel(context.Background())
cliCtx := context.WithValue(backgroundCtx, "gin", c)
var cliClient client.Client
defer func() {
// Ensure the client's mutex is unlocked on function exit.
@@ -327,21 +329,26 @@ outLoop:
// Send the message and receive response chunks and errors via channels.
respChan, errChan := cliClient.SendRawMessageStream(cliCtx, rawJSON, alt)
apiResponseData := make([]byte, 0)
for {
select {
// Handle client disconnection.
case <-c.Request.Context().Done():
if c.Request.Context().Err().Error() == "context canceled" {
log.Debugf("GeminiClient disconnected: %v", c.Request.Context().Err())
c.Set("API_RESPONSE", apiResponseData)
cliCancel() // Cancel the backend request.
return
}
// Process incoming response chunks.
case chunk, okStream := <-respChan:
if !okStream {
c.Set("API_RESPONSE", apiResponseData)
cliCancel()
return
}
apiResponseData = append(apiResponseData, chunk...)
if cliClient.(*client.GeminiClient).GetGenerativeLanguageAPIKey() == "" {
if alt == "" {
responseResult := gjson.GetBytes(chunk, "response")
@@ -382,6 +389,7 @@ outLoop:
c.Status(err.StatusCode)
_, _ = fmt.Fprint(c.Writer, err.Error.Error())
flusher.Flush()
c.Set("API_RESPONSE", []byte(err.Error.Error()))
cliCancel()
}
return
@@ -400,7 +408,9 @@ func (h *GeminiAPIHandlers) handleGeminiCountTokens(c *gin.Context, rawJSON []by
// orgrawJSON := rawJSON
modelResult := gjson.GetBytes(rawJSON, "model")
modelName := modelResult.String()
cliCtx, cliCancel := context.WithCancel(context.Background())
backgroundCtx, cliCancel := context.WithCancel(context.Background())
cliCtx := context.WithValue(backgroundCtx, "gin", c)
var cliClient client.Client
defer func() {
if cliClient != nil {
@@ -441,6 +451,7 @@ func (h *GeminiAPIHandlers) handleGeminiCountTokens(c *gin.Context, rawJSON []by
} else {
c.Status(err.StatusCode)
_, _ = c.Writer.Write([]byte(err.Error.Error()))
c.Set("API_RESPONSE", []byte(err.Error.Error()))
cliCancel()
// log.Debugf(err.Error.Error())
// log.Debugf(string(rawJSON))
@@ -455,6 +466,7 @@ func (h *GeminiAPIHandlers) handleGeminiCountTokens(c *gin.Context, rawJSON []by
}
}
_, _ = c.Writer.Write(resp)
c.Set("API_RESPONSE", resp)
cliCancel()
break
}
@@ -468,7 +480,9 @@ func (h *GeminiAPIHandlers) handleGeminiGenerateContent(c *gin.Context, rawJSON
modelResult := gjson.GetBytes(rawJSON, "model")
modelName := modelResult.String()
cliCtx, cliCancel := context.WithCancel(context.Background())
backgroundCtx, cliCancel := context.WithCancel(context.Background())
cliCtx := context.WithValue(backgroundCtx, "gin", c)
var cliClient client.Client
defer func() {
if cliClient != nil {
@@ -529,6 +543,7 @@ func (h *GeminiAPIHandlers) handleGeminiGenerateContent(c *gin.Context, rawJSON
} else {
c.Status(err.StatusCode)
_, _ = c.Writer.Write([]byte(err.Error.Error()))
c.Set("API_RESPONSE", []byte(err.Error.Error()))
cliCancel()
}
break
@@ -540,6 +555,7 @@ func (h *GeminiAPIHandlers) handleGeminiGenerateContent(c *gin.Context, rawJSON
}
}
_, _ = c.Writer.Write(resp)
c.Set("API_RESPONSE", resp)
cliCancel()
break
}
@@ -570,7 +586,9 @@ func (h *GeminiAPIHandlers) handleCodexStreamGenerateContent(c *gin.Context, raw
modelName := gjson.GetBytes(rawJSON, "model")
cliCtx, cliCancel := context.WithCancel(context.Background())
backgroundCtx, cliCancel := context.WithCancel(context.Background())
cliCtx := context.WithValue(backgroundCtx, "gin", c)
var cliClient client.Client
defer func() {
// Ensure the client's mutex is unlocked on function exit.
@@ -595,6 +613,9 @@ outLoop:
// Send the message and receive response chunks and errors via channels.
respChan, errChan := cliClient.SendRawMessageStream(cliCtx, []byte(newRequestJSON), "")
apiResponseData := make([]byte, 0)
params := &translatorGeminiToCodex.ConvertCodexResponseToGeminiParams{
Model: modelName.String(),
CreatedAt: 0,
@@ -607,15 +628,18 @@ outLoop:
case <-c.Request.Context().Done():
if c.Request.Context().Err().Error() == "context canceled" {
log.Debugf("CodexClient disconnected: %v", c.Request.Context().Err())
c.Set("API_RESPONSE", apiResponseData)
cliCancel() // Cancel the backend request.
return
}
// Process incoming response chunks.
case chunk, okStream := <-respChan:
if !okStream {
c.Set("API_RESPONSE", apiResponseData)
cliCancel()
return
}
apiResponseData = append(apiResponseData, chunk...)
if bytes.HasPrefix(chunk, []byte("data: ")) {
jsonData := chunk[6:]
@@ -643,6 +667,7 @@ outLoop:
c.Status(err.StatusCode)
_, _ = fmt.Fprint(c.Writer, err.Error.Error())
flusher.Flush()
c.Set("API_RESPONSE", []byte(err.Error.Error()))
cliCancel()
}
return
@@ -663,7 +688,9 @@ func (h *GeminiAPIHandlers) handleCodexGenerateContent(c *gin.Context, rawJSON [
modelName := gjson.GetBytes(rawJSON, "model")
cliCtx, cliCancel := context.WithCancel(context.Background())
backgroundCtx, cliCancel := context.WithCancel(context.Background())
cliCtx := context.WithValue(backgroundCtx, "gin", c)
var cliClient client.Client
defer func() {
// Ensure the client's mutex is unlocked on function exit.
@@ -687,21 +714,25 @@ outLoop:
// Send the message and receive response chunks and errors via channels.
respChan, errChan := cliClient.SendRawMessageStream(cliCtx, []byte(newRequestJSON), "")
apiResponseData := make([]byte, 0)
for {
select {
// Handle client disconnection.
case <-c.Request.Context().Done():
if c.Request.Context().Err().Error() == "context canceled" {
log.Debugf("CodexClient disconnected: %v", c.Request.Context().Err())
c.Set("API_RESPONSE", apiResponseData)
cliCancel() // Cancel the backend request.
return
}
// Process incoming response chunks.
case chunk, okStream := <-respChan:
if !okStream {
c.Set("API_RESPONSE", apiResponseData)
cliCancel()
return
}
apiResponseData = append(apiResponseData, chunk...)
if bytes.HasPrefix(chunk, []byte("data: ")) {
jsonData := chunk[6:]
@@ -723,6 +754,7 @@ outLoop:
} else {
c.Status(err.StatusCode)
_, _ = fmt.Fprint(c.Writer, err.Error.Error())
c.Set("API_RESPONSE", []byte(err.Error.Error()))
cliCancel()
}
return