Fix tool schema conversion for Gemini API

The Gemini model was failing to respect 'required' parameter constraints in tool definitions, leading to invalid tool calls.

The root cause was that the proxy was sending a standard JSON Schema to the Gemini API, but the API expects a proprietary, non-standard format (uppercase types, no '$schema' or 'additionalProperties' keywords). This caused the API to return 400 errors or silently ignore the schema.

This commit refactors the tool schema conversion to use strongly-typed structs that represent the Gemini-specific format. A new conversion function recursively builds the schema, ensuring only supported fields are sent and types are correctly formatted.

This also corrects the JSON key for the parameters object from 'parametersJsonSchema' to 'parameters' to match the target API.

Additionally, this commit includes several robustness improvements discovered during debugging:
- Enable TCP KeepAlive on the upstream HTTP client to prevent idle timeouts.
- Add panic-safe checks when parsing the Gemini SSE stream.
- Add a keep-alive pinger to the client-facing SSE stream.

Author: dvcrn

internal/gemini/types.go

@@ -24,8 +24,15 @@ type SystemInstruction struct {
 	Parts []ContentPart `json:"parts,omitempty"`
 }
 
-// JSONSchema represents a JSON schema.
-type JSONSchema map[string]interface{}
+// GeminiParameterSchema defines the proprietary schema format for Gemini function parameters.
+type GeminiParameterSchema struct {
+	Type        string                           `json:"type,omitempty"`
+	Description string                           `json:"description,omitempty"`
+	Properties  map[string]*GeminiParameterSchema `json:"properties,omitempty"`
+	Items       *GeminiParameterSchema           `json:"items,omitempty"`
+	Required    []string                         `json:"required,omitempty"`
+	Enum        []string                         `json:"enum,omitempty"`
+}
 
 // FunctionCall represents a tool call emitted by the model.
 type FunctionCall struct {
@@ -41,9 +48,9 @@ type FunctionResponse struct {
 
 // FunctionDeclaration defines a function that can be called by the model.
 type FunctionDeclaration struct {
-	Name                 string     `json:"name,omitempty"`
-	Description          string     `json:"description,omitempty"`
-	ParametersJsonSchema JSONSchema `json:"parametersJsonSchema,omitempty"`
+	Name        string                 `json:"name,omitempty"`
+	Description string                 `json:"description,omitempty"`
+	Parameters  *GeminiParameterSchema `json:"parameters,omitempty"`
 }
 
 // UnmarshalJSON: accept parametersJsonSchema (camelCase) and parameters (snake_case).
@@ -54,7 +61,7 @@ func (f *FunctionDeclaration) UnmarshalJSON(b []byte) error {
 	var a alias
 	if err := json.Unmarshal(b, &a); err == nil {
 		*f = FunctionDeclaration(a)
-		if f.ParametersJsonSchema != nil {
+		if f.Parameters != nil {
 			return nil
 		}
 	}
@@ -70,17 +77,17 @@ func (f *FunctionDeclaration) UnmarshalJSON(b []byte) error {
 	}
 	// snake_case key used by generativelanguage public API
 	if v, ok := raw["parameters"]; ok {
-		var m map[string]interface{}
-		if err := json.Unmarshal(v, &m); err == nil {
-			f.ParametersJsonSchema = m
+		var schema GeminiParameterSchema
+		if err := json.Unmarshal(v, &schema); err == nil {
+			f.Parameters = &schema
 			return nil
 		}
 	}
 	// or explicit camelCase if present but empty earlier
-	if v, ok := raw["parametersJsonSchema"]; ok && f.ParametersJsonSchema == nil {
-		var m map[string]interface{}
-		if err := json.Unmarshal(v, &m); err == nil {
-			f.ParametersJsonSchema = m
+	if v, ok := raw["parametersJsonSchema"]; ok && f.Parameters == nil {
+		var schema GeminiParameterSchema
+		if err := json.Unmarshal(v, &schema); err == nil {
+			f.Parameters = &schema
 		}
 	}
 	return nil

internal/http/http_client_default.go

@@ -3,6 +3,7 @@
 package http
 
 import (
+	"net"
 	"net/http"
 	"time"
 )
@@ -11,6 +12,10 @@ import (
 func NewHTTPClient() HTTPClient {
 	return &http.Client{
 		Transport: &http.Transport{
+			DialContext: (&net.Dialer{
+				Timeout:   30 * time.Second,
+				KeepAlive: 30 * time.Second,
+			}).DialContext,
 			MaxIdleConns:        100,
 			MaxIdleConnsPerHost: 10,
 			IdleConnTimeout:     90 * time.Second,

internal/logger/logger.go

@@ -44,6 +44,18 @@ func colorize(s interface{}, c int) string {
 func newLogger() *zerolog.Logger {
 	env := os.Getenv("ENV")
 
+	// Set log level based on LOG_LEVEL env var, default to info
+	logLevel := zerolog.InfoLevel
+	if levelStr := os.Getenv("LOG_LEVEL"); levelStr != "" {
+		if parsedLevel, err := zerolog.ParseLevel(strings.ToLower(levelStr)); err == nil {
+			logLevel = parsedLevel
+		} else {
+			fmt.Fprintf(os.Stderr, "Invalid LOG_LEVEL \"%s\"; defaulting to 'info'\n", levelStr)
+		}
+	}
+
+	zerolog.SetGlobalLevel(logLevel)
+
 	if env == "development" || env == "dev" || env == "" {
 		return newDevelopment()
 	}

internal/openai/stream_transformer.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/dvcrn/gemini-code-assist-proxy/internal/logger"
 	"github.com/google/uuid"
 )
 
@@ -128,6 +129,8 @@ func CreateOpenAIStreamTransformer(model string) func(<-chan StreamChunk) <-chan
 
 			// Process each chunk
 			for chunk := range input {
+				logger.Get().Debug().Interface("chunk", chunk).Msg("Processing Gemini stream chunk")
+
 				delta := OpenAIDelta{}
 				shouldSend := false
 
@@ -222,7 +225,9 @@ func CreateOpenAIStreamTransformer(model string) func(<-chan StreamChunk) <-chan
 					}
 
 					if jsonBytes, err := json.Marshal(openAIChunk); err == nil {
-						output <- fmt.Sprintf("data: %s\n\n", string(jsonBytes))
+						sse := fmt.Sprintf("data: %s\n\n", string(jsonBytes))
+						logger.Get().Debug().Str("sse", sse).Msg("Sending OpenAI SSE chunk")
+						output <- sse
 					}
 				}
 			}

internal/server/chat_completions_handler.go

@@ -1,6 +1,7 @@
 package server
 
 import (
+	"context"
 	"encoding/json"
 	"io"
 	"net/http"
@@ -147,6 +148,31 @@ func (s *Server) chatCompletionRequestStream(w http.ResponseWriter, r *http.Requ
 	logger.Get().Info().
 		Str("model", gemReq.Model).
 		Msg("Starting upstream StreamGenerateContent")
+
+	// Pinger to keep connection alive
+	pingerCtx, cancelPinger := context.WithCancel(r.Context())
+	defer cancelPinger()
+
+	go func() {
+		ticker := time.NewTicker(10 * time.Second)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-ticker.C:
+				logger.Get().Debug().Msg("Sending SSE ping to keep connection alive")
+				if _, err := io.WriteString(w, ": ping\n\n"); err != nil {
+					logger.Get().Warn().Err(err).Msg("Failed to write SSE ping")
+					return
+				}
+				if flusher, ok := w.(http.Flusher); ok {
+					flusher.Flush()
+				}
+			case <-pingerCtx.Done():
+				return
+			}
+		}
+	}()
+
 	if err := s.geminiClient.StreamGenerateContent(r.Context(), gemReq, upstream); err != nil {
 		logger.Get().Error().Err(err).Msg("StreamGenerateContent call failed")
 		http.Error(w, "Upstream streaming error", http.StatusInternalServerError)
@@ -161,6 +187,9 @@ func (s *Server) chatCompletionRequestStream(w http.ResponseWriter, r *http.Requ
 		firstUpstream := true
 		firstThoughtSeen := false
 		for line := range upstream {
+			if firstUpstream {
+				cancelPinger() // Stop pinger on first data
+			}
 			// Process only data lines
 			if !strings.HasPrefix(line, "data: ") {
 				continue
@@ -207,7 +236,11 @@ func (s *Server) chatCompletionRequestStream(w http.ResponseWriter, r *http.Requ
 			// Extract candidate content parts
 			if cands, ok := obj["candidates"].([]interface{}); ok {
 				for _, c := range cands {
-					cand, _ := c.(map[string]interface{})
+					cand, ok := c.(map[string]interface{})
+					if !ok {
+						logger.Get().Warn().Interface("candidate", c).Msg("Skipping invalid candidate in Gemini stream")
+						continue
+					}
 
 					// Optional grounding metadata passthrough
 					if gm, ok := cand["groundingMetadata"]; ok && gm != nil {
@@ -229,7 +262,11 @@ func (s *Server) chatCompletionRequestStream(w http.ResponseWriter, r *http.Requ
 
 					// Process parts
 					for _, p := range parts {
-						part, _ := p.(map[string]interface{})
+						part, ok := p.(map[string]interface{})
+						if !ok {
+							logger.Get().Warn().Interface("part", p).Msg("Skipping invalid part in Gemini stream")
+							continue
+						}
 
 						// Thought tokens (reasoning) — map to OpenAI reasoning stream
 						if isThought, ok := part["thought"].(bool); ok && isThought {

internal/transform/openai_to_gemini.go

@@ -248,16 +248,29 @@ func convertToolsToGeminiTools(tools []openai.Tool) []gemini.Tool {
 			continue
 		}
 
-		var schema gemini.JSONSchema
+		var geminiSchema *gemini.GeminiParameterSchema
 		if m, ok := t.Function.Parameters.(map[string]interface{}); ok {
-			schema = m
+			geminiSchema = convertToGeminiSchema(m)
 		}
 
-		fns = append(fns, gemini.FunctionDeclaration{
-			Name:                 t.Function.Name,
-			Description:          t.Function.Description,
-			ParametersJsonSchema: schema,
-		})
+		convertedFn := gemini.FunctionDeclaration{
+			Name:        t.Function.Name,
+			Description: t.Function.Description,
+			Parameters:  geminiSchema,
+		}
+
+		// For specific tools, log the before and after transformation for debugging
+		if t.Function.Name == "Read" || t.Function.Name == "Edit" || t.Function.Name == "MultiEdit" {
+			originalJSON, _ := json.Marshal(t.Function)
+			convertedJSON, _ := json.Marshal(convertedFn)
+			logger.Get().Info().
+				Str("tool_name", t.Function.Name).
+				RawJSON("original_schema", originalJSON).
+				RawJSON("converted_schema", convertedJSON).
+				Msg("Dumping tool schema conversion from OpenAI to Gemini")
+		}
+
+		fns = append(fns, convertedFn)
 	}
 
 	if len(fns) == 0 {
@@ -268,3 +281,51 @@ func convertToolsToGeminiTools(tools []openai.Tool) []gemini.Tool {
 		{FunctionDeclarations: fns},
 	}
 }
+
+// convertToGeminiSchema recursively converts a generic map representing a JSON schema
+// into the strongly-typed GeminiParameterSchema struct, only mapping supported fields.
+func convertToGeminiSchema(input map[string]interface{}) *gemini.GeminiParameterSchema {
+	if input == nil {
+		return nil
+	}
+
+	output := &gemini.GeminiParameterSchema{}
+
+	if t, ok := input["type"].(string); ok {
+		output.Type = strings.ToUpper(t)
+	}
+	if d, ok := input["description"].(string); ok {
+		output.Description = d
+	}
+
+	if r, ok := input["required"].([]interface{}); ok {
+		for _, v := range r {
+			if s, ok := v.(string); ok {
+				output.Required = append(output.Required, s)
+			}
+		}
+	}
+
+	if e, ok := input["enum"].([]interface{}); ok {
+		for _, v := range e {
+			if s, ok := v.(string); ok {
+				output.Enum = append(output.Enum, s)
+			}
+		}
+	}
+
+	if p, ok := input["properties"].(map[string]interface{}); ok {
+		output.Properties = make(map[string]*gemini.GeminiParameterSchema)
+		for k, v := range p {
+			if vMap, ok := v.(map[string]interface{}); ok {
+				output.Properties[k] = convertToGeminiSchema(vMap)
+			}
+		}
+	}
+
+	if i, ok := input["items"].(map[string]interface{}); ok {
+		output.Items = convertToGeminiSchema(i)
+	}
+
+	return output
+}