feat: implement LLM robustness features including text-to-tool-call recovery, empty response nudges, output truncation, and API timeouts.

Author: ngoyal16

adk/connection/local_connection.go

@@ -230,7 +230,20 @@ func NewLocalConnection(ctx context.Context, cfg LocalConfig, logger *slog.Logge
 		return nil, fmt.Errorf("unexpected message type received during initialization: %T", initResp.Payload)
 	}
 
-	// 10. Launch background read loop
+	// 10. Configure WebSocket keepalive.
+	// The server sends periodic pings; gorilla/websocket auto-responds with pongs.
+	// Set a read deadline that resets on each ping to detect server death.
+	// 45s pong timeout matches the server's wsPongTimeout constant.
+	const wsPongTimeout = 45 * time.Second
+	wsConn.SetReadDeadline(time.Now().Add(wsPongTimeout))
+	wsConn.SetPingHandler(func(appData string) error {
+		// Reset read deadline on each ping from server
+		wsConn.SetReadDeadline(time.Now().Add(wsPongTimeout))
+		// Send pong response (gorilla requires explicit handler to reset deadline)
+		return wsConn.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(5*time.Second))
+	})
+
+	// 11. Launch background read loop
 	go c.readLoop(ctx)
 
 	return c, nil
@@ -452,11 +465,21 @@ func (c *LocalConnection) handleStepUpdate(su *pb.StepUpdate) {
 }
 
 func (c *LocalConnection) handleTrajectoryState(ts *pb.TrajectoryState) {
+	c.logger.Debug("received trajectory state",
+		"state", ts.State.String(),
+		"parent", ts.ParentTrajectoryId,
+		"trajectory_id", ts.TrajectoryId,
+	)
 	// Root trajectory state change to idle/done/error signals the end of the Chat turn.
 	if ts.ParentTrajectoryId == "" {
 		if ts.State == pb.TrajectoryState_TRAJ_IDLE ||
 			ts.State == pb.TrajectoryState_TRAJ_COMPLETED ||
 			ts.State == pb.TrajectoryState_TRAJ_ERROR {
+			if ts.State == pb.TrajectoryState_TRAJ_ERROR {
+				c.logger.Warn("trajectory ended with ERROR state",
+					"trajectory_id", ts.TrajectoryId,
+				)
+			}
 			c.closeStepsCh()
 		}
 	}

agents/zenith/config.go

@@ -159,10 +159,11 @@ func (z *ZenithConfig) ApplyTo(cfg *adk.LocalAgentConfig) error {
 			cfProvider = "workers-ai" // default to Workers AI
 		}
 		cfg.CloudflareProvider = cfProvider
+		cfg.CloudflareCacheTTL = 3600 // Cache responses for 1 hour
 		if z.Model != "" {
 			cfg.CloudflareModel = z.Model
 		} else {
-			cfg.CloudflareModel = "@cf/zai-org/glm-4.7-flash" // sensible default
+			cfg.CloudflareModel = "@cf/google/gemma-4-26b-a4b-it" // best free model for agentic work
 		}
 
 	default:

internal/engine/engine.go

@@ -80,6 +80,7 @@ type Engine struct {
 	msgCtx              MessageContextConfig         // Per-message context enrichment config
 	notifyCh            <-chan tools.SystemMessage    // System notifications (timers, task completions)
 	running             atomic.Int32                 // 1 = engine is running a turn, 0 = idle
+	preCompletionHook   func()                       // Called before TRAJ_IDLE so session can save state
 
 	// Subagent support
 	depth              int           // Nesting depth (0 = root)
@@ -337,6 +338,14 @@ func (e *Engine) SetSettingsChanges(changes []SettingsChange) {
 	e.msgCtx.SettingsChanges = append(e.msgCtx.SettingsChanges, changes...)
 }
 
+// SetPreCompletionHook registers a callback that runs after the engine's
+// final response but BEFORE TRAJ_IDLE is emitted over the WebSocket.
+// The session uses this to save conversation state before the SDK can
+// kill the harness process in response to TRAJ_IDLE.
+func (e *Engine) SetPreCompletionHook(fn func()) {
+	e.preCompletionHook = fn
+}
+
 // RunWithContext executes the agentic loop with optional per-message host context.
 // The host context (active file, cursor, etc.) is injected into the enriched message
 // sent to the LLM, while the raw userMessage is preserved for step updates.
@@ -398,6 +407,7 @@ drained:
 	toolDecls := e.buildToolDeclarations()
 
 	// Agentic loop
+	emptyRetries := 0
 	for turn := 0; turn < e.maxTurns; turn++ {
 		select {
 		case <-ctx.Done():
@@ -478,11 +488,16 @@ drained:
 		var err error
 		start := time.Now()
 
+		// Per-call timeout prevents hanging when the LLM API stalls
+		// (e.g., Cloudflare worker dies mid-stream without closing the connection).
+		callCtx, callCancel := context.WithTimeout(ctx, llmCallTimeout)
+
 		if sp, ok := e.provider.(llm.StreamingProvider); ok {
-			resp, err = e.streamGenerate(ctx, sp, req)
+			resp, err = e.streamGenerate(callCtx, sp, req)
 		} else {
-			resp, err = e.provider.Generate(ctx, req)
+			resp, err = e.provider.Generate(callCtx, req)
 		}
+		callCancel()
 		latency := time.Since(start)
 
 		// Trace the response
@@ -529,6 +544,43 @@ drained:
 			continue // Retry the turn
 		}
 
+		// ── Text-to-tool-call recovery ──
+		// Some models (e.g., Llama 3.3 via Workers AI) output tool calls as text
+		// instead of structured tool_calls. Detect and recover so the agentic
+		// loop continues instead of stalling.
+		if len(resp.ToolCalls) == 0 && resp.Content != "" {
+			recovered, remaining := tryExtractToolCallsFromText(resp.Content, e.knownToolNames(), e.logger)
+			if len(recovered) > 0 {
+				resp.ToolCalls = recovered
+				resp.Content = remaining
+				resp.FinishReason = "tool_calls"
+			}
+		}
+
+		// Empty model response recovery — model returned no text AND no tool calls
+		// with a non-max_tokens finish reason. This happens with smaller models
+		// (Workers AI, Ollama) that occasionally produce empty completions.
+		if resp.Content == "" && len(resp.ToolCalls) == 0 && resp.FinishReason != "max_tokens" {
+			emptyRetries++
+			if emptyRetries <= 2 {
+				e.logger.Warn("model returned empty response, retrying with nudge",
+					"finish_reason", resp.FinishReason,
+					"retry", emptyRetries,
+					"turn", turn,
+				)
+				e.history = append(e.history, llm.Message{
+					Role:    "model",
+					Content: "[Empty response]",
+				})
+				e.history = append(e.history, llm.Message{
+					Role:    "user",
+					Content: "Your previous response was empty. Please continue with the task. If you need to use a tool, call it. If you're done, provide your final answer.",
+				})
+				continue
+			}
+			// After 2 retries, fall through to handleFinalResponse with empty text
+		}
+
 		// If model returned text with no tool calls → final response
 		if resp.FinishReason == "stop" || len(resp.ToolCalls) == 0 {
 			return e.handleFinalResponse(resp)
@@ -567,10 +619,28 @@ drained:
 	return fmt.Errorf("engine: exceeded max turns (%d)", e.maxTurns)
 }
 
+// llmCallTimeout is the maximum time allowed for a single LLM API call
+// (including streaming). Prevents indefinite hangs when the API stalls.
+const llmCallTimeout = 120 * time.Second
+
+// maxToolResultSize is the maximum size (in bytes) of a single tool result
+// added to the conversation history. Results exceeding this are truncated to
+// prevent overwhelming smaller LLMs (e.g., Workers AI free-tier models).
+// 32KB ≈ 8K tokens — enough for meaningful content while staying within
+// context budget limits for most models.
+const maxToolResultSize = 32_000
+
 // toolResultMsg builds a tool result message, propagating the ThoughtSignature
 // from the ToolCall. Gemini 3.5+ requires thought_signature on functionResponse
 // parts to maintain chain integrity.
 func toolResultMsg(tc llm.ToolCall, content string, isError bool) llm.Message {
+	// Truncate oversized tool results to prevent context blowup
+	if len(content) > maxToolResultSize {
+		content = content[:maxToolResultSize] + fmt.Sprintf(
+			"\n\n... [output truncated, showing %d/%d bytes. Use more specific queries or line ranges to get targeted results.]",
+			maxToolResultSize, len(content),
+		)
+	}
 	return llm.Message{
 		Role: "tool",
 		ToolResult: &llm.ToolCallResult{
@@ -660,6 +730,13 @@ func (e *Engine) handleFinalResponse(resp *llm.GenerateResponse) error {
 		Content: resp.Content,
 	})
 
+	// Run pre-completion hook (e.g., session saves conversation state)
+	// BEFORE emitting TRAJ_IDLE, because the SDK will kill the process
+	// as soon as it receives TRAJ_IDLE.
+	if e.preCompletionHook != nil {
+		e.preCompletionHook()
+	}
+
 	e.emitTrajectoryState(pb.TrajectoryState_TRAJ_IDLE)
 	return nil
 }
@@ -1468,7 +1545,10 @@ func (e *Engine) extractToolResult(step *pb.StepUpdate) string {
 				"size_bytes": e.SizeBytes, "child_count": e.ChildCount,
 			})
 		}
-		result = map[string]interface{}{"entries": entries}
+		result = map[string]interface{}{
+			"entries":       entries,
+			"total_entries": len(a.ListDir.Entries),
+		}
 	case *pb.StepUpdate_GrepSearch:
 		var matches []map[string]interface{}
 		for _, m := range a.GrepSearch.Matches {
@@ -1701,3 +1781,44 @@ func (e *Engine) buildToolDeclarations() []llm.FunctionDeclaration {
 func (e *Engine) History() []llm.Message {
 	return e.history
 }
+
+// knownToolNames returns a set of all tool names the engine knows about.
+// This includes built-in tools, engine-intercepted tools, host tools, and MCP tools.
+// Used by text-to-tool-call recovery to validate extracted tool names.
+func (e *Engine) knownToolNames() map[string]bool {
+	names := make(map[string]bool)
+
+	// Built-in + engine-intercepted tools from the registry
+	for _, s := range e.toolRegistry.Schemas() {
+		names[s.Name] = true
+	}
+
+	// Host-side tools
+	for name := range e.hostToolNames {
+		names[name] = true
+	}
+
+	// Host tool declarations (some may not be in hostToolNames)
+	for _, d := range e.hostToolDecls {
+		names[d.Name] = true
+	}
+
+	// MCP tools
+	if e.mcpMgr != nil {
+		for _, d := range e.mcpMgr.ToolDeclarations() {
+			names[d.Name] = true
+		}
+	}
+
+	// Engine-intercepted tools (not in registry)
+	for _, name := range []string{
+		"invoke_subagent", "define_subagent", "manage_subagents",
+		"send_message", "ask_question", "ask_permission", "list_permissions",
+		"knowledge_write", "knowledge_replace", "knowledge_delete",
+		"publish", "finish",
+	} {
+		names[name] = true
+	}
+
+	return names
+}