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