diff --git a/.gitignore b/.gitignore index 0fa0f42..33280d8 100644 --- a/.gitignore +++ b/.gitignore @@ -18,3 +18,4 @@ coverage.html .DS_Store Thumbs.db LOAD_TEST_REPORT.md +.alma-snapshots diff --git a/internal/openai/responses_compat.go b/internal/openai/responses_compat.go index 8037be7..b3d6604 100644 --- a/internal/openai/responses_compat.go +++ b/internal/openai/responses_compat.go @@ -56,8 +56,30 @@ type StreamEvent struct { // PrepareCompatibleResponses maps current OpenAI/Codex Responses request // shapes to the subset understood by the Grok CLI upstream. func PrepareCompatibleResponses(body map[string]any) (map[string]any, *ResponsesCompatibility, error) { + return PrepareCompatibleResponsesWithCache(body, DefaultToolReplay) +} + +// PrepareCompatibleResponsesWithCache is the testable entry point that accepts +// an explicit tool-call replay cache (Alma multi-turn continuity). +// +// Pipeline (Alma/Codex multi-turn tool continuity): +// 1. expand item_reference from cache +// 2. re-insert missing function/custom calls for tool outputs (prev-resp / cache) +// 3. normalize each input item (ModelInput hygiene, custom→function, drop residual refs) +// 4. prune remaining orphan tool outputs +// +// Note: normalize runs after replay insert so cached custom_tool_call items +// still go through this package's alias rewrite. Cached function_call items +// are already minimal ModelInput shapes. +func PrepareCompatibleResponsesWithCache(body map[string]any, cache *ToolReplayCache) (map[string]any, *ResponsesCompatibility, error) { out := PrepareResponses(body) removeEncryptedReasoningInclude(out) + // Grok CLI has no previous_response_id; affinity/replay consume it first. + previousResponseID := String(body, "previous_response_id", "") + promptCacheKey := String(body, "prompt_cache_key", "") + delete(out, "previous_response_id") + model := String(body, "model", "") + compat := &ResponsesCompatibility{ aliases: make(map[string]toolIdentity), originalAliases: make(map[string]string), @@ -72,18 +94,18 @@ func PrepareCompatibleResponses(body map[string]any) (map[string]any, *Responses } if input, ok := out["input"].([]any); ok { - rewritten := make([]any, 0, len(input)) - for index, raw := range input { - item, extra, err := compat.normalizeInputItem(raw, index) - if err != nil { - return nil, nil, err - } - sources = append(sources, extra...) - if item != nil { - rewritten = append(rewritten, item) - } + input = expandItemReferences(cache, model, input) + probe := map[string]any{"input": input} + applyToolCallReplay(cache, model, probe, previousResponseID, promptCacheKey) + input, _ = probe["input"].([]any) + + rewritten, extra, err := compat.normalizeInputList(input) + if err != nil { + return nil, nil, err } + sources = append(sources, extra...) out["input"] = rewritten + pruneOrphanToolOutputs(out) } tools, err := compat.normalizeToolSources(sources) @@ -97,6 +119,22 @@ func PrepareCompatibleResponses(body map[string]any) (map[string]any, *Responses return out, compat, nil } +func (c *ResponsesCompatibility) normalizeInputList(input []any) ([]any, []toolSource, error) { + rewritten := make([]any, 0, len(input)) + sources := make([]toolSource, 0) + for index, raw := range input { + item, extra, err := c.normalizeInputItem(raw, index) + if err != nil { + return nil, nil, err + } + sources = append(sources, extra...) + if item != nil { + rewritten = append(rewritten, item) + } + } + return rewritten, sources, nil +} + func (c *ResponsesCompatibility) normalizeInputItem(raw any, index int) (any, []toolSource, error) { item, ok := raw.(map[string]any) if !ok { @@ -107,11 +145,31 @@ func (c *ResponsesCompatibility) normalizeInputItem(raw any, index int) (any, [] return item, nil, nil } switch kind { - case "message", "function_call_output", "web_search_call": + case "message": return sanitizeInputItem(item), nil, nil + case "function_call_output": + out := sanitizeInputItem(item) + out["output"] = flattenToolOutput(item["output"]) + return out, nil, nil + // Server-tool history is not part of Grok CLI ModelInput. Codex re-sends + // visible search results as messages when needed; forwarding these items + // yields a hard upstream 422 (untagged enum ModelInput). + case "web_search_call", "x_search_call", "web_search", "x_search", + "computer_call", "computer_call_output", + "local_shell_call_output", "image_generation_call", + "code_interpreter_call", "file_search_call", + "mcp_call", "mcp_list_tools", "mcp_approval_request", "mcp_approval_response": + return nil, nil, nil case "reasoning": + // Codex/Alma often re-send reasoning shells with content:null / + // encrypted_content:null. Null fields fail Grok's untagged ModelInput + // enum (422), so strip nulls. Foreign encrypted_content is never + // portable across the account pool. out := sanitizeInputItem(item) delete(out, "encrypted_content") + if out["content"] == nil { + delete(out, "content") + } if !hasReasoningText(out) { return nil, nil, nil } @@ -128,8 +186,14 @@ func (c *ResponsesCompatibility) normalizeInputItem(raw any, index int) (any, [] } alias := c.alias(identity) out["name"] = alias + // Grok ModelInput function_call.arguments must be a JSON string. + out["arguments"] = stringifyToolArguments(item["arguments"]) delete(out, "namespace") return out, nil, nil + case "item_reference": + // Expanded earlier from the replay cache. Any leftover reference is not + // a Grok ModelInput variant and must not be forwarded. + return nil, nil, nil case "additional_tools": return nil, toolSources(item["tools"], true), nil case "tool_search_call": @@ -155,6 +219,7 @@ func (c *ResponsesCompatibility) normalizeInputItem(raw any, index int) (any, [] case "custom_tool_call_output": out := sanitizeInputItem(item) out["type"] = "function_call_output" + out["output"] = flattenToolOutput(item["output"]) delete(out, "name") return out, nil, nil case "agent_message": @@ -510,6 +575,66 @@ func encodeCustomInput(value any) string { return string(encoded) } +// stringifyToolArguments forces Grok ModelInput function_call.arguments to a +// JSON string. Some clients send structured objects. +func stringifyToolArguments(value any) string { + switch typed := value.(type) { + case nil: + return "{}" + case string: + if strings.TrimSpace(typed) == "" { + return "{}" + } + return typed + default: + encoded, err := json.Marshal(typed) + if err != nil { + return "{}" + } + return string(encoded) + } +} + +// flattenToolOutput collapses array/object tool outputs into a single string +// accepted by the Grok CLI function_call_output shape. +func flattenToolOutput(value any) string { + switch typed := value.(type) { + case nil: + return "" + case string: + return typed + case []any: + parts := make([]string, 0, len(typed)) + for _, entry := range typed { + switch part := entry.(type) { + case string: + parts = append(parts, part) + case map[string]any: + if text := String(part, "text", ""); text != "" { + parts = append(parts, text) + continue + } + encoded, err := json.Marshal(part) + if err == nil { + parts = append(parts, string(encoded)) + } + default: + encoded, err := json.Marshal(part) + if err == nil { + parts = append(parts, string(encoded)) + } + } + } + return strings.Join(parts, "\n") + default: + encoded, err := json.Marshal(typed) + if err != nil { + return fmt.Sprint(typed) + } + return string(encoded) + } +} + func decodeCustomInput(value any) string { text, _ := value.(string) var wrapper map[string]any diff --git a/internal/openai/tool_replay.go b/internal/openai/tool_replay.go new file mode 100644 index 0000000..3a1444b --- /dev/null +++ b/internal/openai/tool_replay.go @@ -0,0 +1,474 @@ +package openai + +import ( + "strings" + "sync" + "time" +) + +// ToolReplayCache stores completed function/custom tool calls so multi-turn +// OpenAI Responses clients (notably Alma) can continue with only +// previous_response_id / item_reference + tool outputs. +// +// Only tool calls are cached — never encrypted reasoning — so a multi-account +// pool cannot re-inject another account's opaque ciphertext. +type ToolReplayCache struct { + mu sync.Mutex + ttl time.Duration + max int + entries map[string]toolReplayEntry +} + +type toolReplayEntry struct { + items []map[string]any + expiresAt time.Time +} + +const ( + defaultToolReplayTTL = time.Hour + defaultToolReplayMax = 10240 +) + +// DefaultToolReplay is the process-wide cache used by the Responses path. +var DefaultToolReplay = NewToolReplayCache(defaultToolReplayTTL, defaultToolReplayMax) + +func NewToolReplayCache(ttl time.Duration, maxEntries int) *ToolReplayCache { + if ttl <= 0 { + ttl = defaultToolReplayTTL + } + if maxEntries < 1 { + maxEntries = defaultToolReplayMax + } + return &ToolReplayCache{ + ttl: ttl, + max: maxEntries, + entries: make(map[string]toolReplayEntry), + } +} + +func toolReplayKey(model, key string) string { + return model + "\x00" + key +} + +func (c *ToolReplayCache) Get(model, key string) []map[string]any { + if c == nil || model == "" || key == "" { + return nil + } + c.mu.Lock() + defer c.mu.Unlock() + entry, ok := c.entries[toolReplayKey(model, key)] + if !ok { + return nil + } + if time.Now().After(entry.expiresAt) { + delete(c.entries, toolReplayKey(model, key)) + return nil + } + // Sliding TTL: each successful read refreshes the entry expiry. + entry.expiresAt = time.Now().Add(c.ttl) + c.entries[toolReplayKey(model, key)] = entry + return cloneToolReplayItems(entry.items) +} + +func (c *ToolReplayCache) Put(model, key string, items []map[string]any) { + if c == nil || model == "" || key == "" || len(items) == 0 { + return + } + normalized := normalizeReplayItems(items) + if len(normalized) == 0 { + return + } + c.mu.Lock() + defer c.mu.Unlock() + if len(c.entries) >= c.max { + c.evictExpiredLocked(time.Now()) + } + if len(c.entries) >= c.max { + for existing := range c.entries { + delete(c.entries, existing) + break + } + } + c.entries[toolReplayKey(model, key)] = toolReplayEntry{ + items: cloneToolReplayItems(normalized), + expiresAt: time.Now().Add(c.ttl), + } +} + +func (c *ToolReplayCache) evictExpiredLocked(now time.Time) { + for key, entry := range c.entries { + if now.After(entry.expiresAt) { + delete(c.entries, key) + } + } +} + +func cloneToolReplayItems(items []map[string]any) []map[string]any { + out := make([]map[string]any, 0, len(items)) + for _, item := range items { + out = append(out, clone(item)) + } + return out +} + +// normalizeReplayItems keeps only the minimal function/custom tool-call shape +// that Grok accepts when replaying prior turns. +func normalizeReplayItems(items []map[string]any) []map[string]any { + out := make([]map[string]any, 0, len(items)) + for _, item := range items { + switch String(item, "type", "") { + case "function_call": + callID := strings.TrimSpace(String(item, "call_id", "")) + name := strings.TrimSpace(String(item, "name", "")) + args, ok := item["arguments"].(string) + if callID == "" || name == "" || !ok { + continue + } + normalized := map[string]any{ + "type": "function_call", + "call_id": callID, + "name": name, + "arguments": args, + } + if id := strings.TrimSpace(String(item, "id", "")); id != "" { + normalized["id"] = id + } + if namespace := strings.TrimSpace(String(item, "namespace", "")); namespace != "" { + normalized["namespace"] = namespace + } + out = append(out, normalized) + case "custom_tool_call": + callID := strings.TrimSpace(String(item, "call_id", "")) + name := strings.TrimSpace(String(item, "name", "")) + if callID == "" || name == "" || item["input"] == nil { + continue + } + normalized := map[string]any{ + "type": "custom_tool_call", + "status": "completed", + "call_id": callID, + "name": name, + "input": item["input"], + } + if id := strings.TrimSpace(String(item, "id", "")); id != "" { + normalized["id"] = id + } + if status := strings.TrimSpace(String(item, "status", "")); status != "" { + normalized["status"] = status + } + if namespace := strings.TrimSpace(String(item, "namespace", "")); namespace != "" { + normalized["namespace"] = namespace + } + out = append(out, normalized) + } + } + return out +} + +// RememberCompletedResponse indexes tool calls from a completed Responses +// object. model must be the client-requested model (not free-tier rewrite). +// +// Keys: +// - prev-resp:{response.id} +// - cache:{prompt_cache_key} when present +// - item:{item.id} for each tool call +func RememberCompletedResponse(cache *ToolReplayCache, model string, response map[string]any, promptCacheKey string) { + if cache == nil || response == nil || strings.TrimSpace(model) == "" { + return + } + items := extractReplayToolCalls(response["output"]) + if len(items) == 0 { + return + } + responseID := String(response, "id", "") + if responseID != "" { + cache.Put(model, "prev-resp:"+responseID, items) + } + if promptCacheKey = strings.TrimSpace(promptCacheKey); promptCacheKey != "" { + cache.Put(model, "cache:"+promptCacheKey, items) + } + for _, item := range items { + if id := String(item, "id", ""); id != "" { + cache.Put(model, "item:"+id, []map[string]any{item}) + } + } +} + +// RememberStreamToolCall indexes one tool call from output_item.done. +// Session keys (prev-resp / cache) are committed on response.completed after +// patching empty output; item:{id} is safe to write immediately so +// item_reference works mid-stream. +func RememberStreamToolCall(cache *ToolReplayCache, model string, item map[string]any, responseID, promptCacheKey string) { + if cache == nil || item == nil || strings.TrimSpace(model) == "" { + return + } + items := normalizeReplayItems([]map[string]any{item}) + if len(items) == 0 { + return + } + if responseID = strings.TrimSpace(responseID); responseID != "" { + cache.Put(model, "prev-resp:"+responseID, mergeReplayItems(cache.Get(model, "prev-resp:"+responseID), items)) + } + if promptCacheKey = strings.TrimSpace(promptCacheKey); promptCacheKey != "" { + cache.Put(model, "cache:"+promptCacheKey, mergeReplayItems(cache.Get(model, "cache:"+promptCacheKey), items)) + } + for _, normalized := range items { + if id := String(normalized, "id", ""); id != "" { + cache.Put(model, "item:"+id, []map[string]any{normalized}) + } + } +} + +// mergeReplayItems appends newItems, de-duplicating by call_id (then id). +func mergeReplayItems(existing, newItems []map[string]any) []map[string]any { + if len(existing) == 0 { + return newItems + } + seen := make(map[string]struct{}, len(existing)+len(newItems)) + out := make([]map[string]any, 0, len(existing)+len(newItems)) + add := func(item map[string]any) { + key := strings.TrimSpace(String(item, "call_id", "")) + if key == "" { + key = strings.TrimSpace(String(item, "id", "")) + } + if key != "" { + if _, ok := seen[key]; ok { + return + } + seen[key] = struct{}{} + } + out = append(out, item) + } + for _, item := range existing { + add(item) + } + for _, item := range newItems { + add(item) + } + return out +} + +func extractReplayToolCalls(raw any) []map[string]any { + output, ok := raw.([]any) + if !ok { + return nil + } + items := make([]map[string]any, 0, len(output)) + for _, entry := range output { + item, ok := entry.(map[string]any) + if !ok { + continue + } + switch String(item, "type", "") { + case "function_call", "custom_tool_call": + items = append(items, clone(item)) + } + } + return items +} + +// expandItemReferences replaces item_reference with cached tool calls. +// Unresolved references are left for normalize to drop. +func expandItemReferences(cache *ToolReplayCache, model string, input []any) []any { + if cache == nil || len(input) == 0 { + return input + } + var out []any + for index, raw := range input { + item, ok := raw.(map[string]any) + if !ok || String(item, "type", "") != "item_reference" { + if out != nil { + out = append(out, raw) + } + continue + } + id := strings.TrimSpace(String(item, "id", "")) + if id == "" { + if out != nil { + out = append(out, raw) + } + continue + } + cached := cache.Get(model, "item:"+id) + if len(cached) == 0 { + if out != nil { + out = append(out, raw) + } + continue + } + if out == nil { + out = make([]any, 0, len(input)+len(cached)-1) + out = append(out, input[:index]...) + } + for _, c := range cached { + out = append(out, c) + } + } + if out == nil { + return input + } + return out +} + +// applyToolCallReplay re-inserts cached function/custom calls for matching +// tool outputs from previous_response_id / prompt_cache_key sessions. +func applyToolCallReplay(cache *ToolReplayCache, model string, body map[string]any, previousResponseID, promptCacheKey string) { + if cache == nil { + return + } + input, ok := body["input"].([]any) + if !ok || len(input) == 0 { + return + } + + // Collect existing calls and outputs already present in the request input. + existingCalls := make(map[string]struct{}) + existingOutputs := make(map[string]struct{}) + for _, raw := range input { + item, ok := raw.(map[string]any) + if !ok { + continue + } + callID := strings.TrimSpace(String(item, "call_id", "")) + switch String(item, "type", "") { + case "function_call", "custom_tool_call": + if callID != "" { + existingCalls[callID] = struct{}{} + } + case "function_call_output", "custom_tool_call_output": + if callID != "" { + existingOutputs[callID] = struct{}{} + } + } + } + if len(existingOutputs) == 0 { + return + } + + // Prefer previous_response_id session, then prompt_cache_key session. + var candidates []map[string]any + if previousResponseID != "" { + candidates = append(candidates, cache.Get(model, "prev-resp:"+previousResponseID)...) + } + if promptCacheKey != "" { + candidates = append(candidates, cache.Get(model, "cache:"+promptCacheKey)...) + } + if len(candidates) == 0 { + return + } + + // Filter: only function/custom calls whose call_id has a matching output + // and is not already present in input. + filtered := make([]map[string]any, 0) + seen := make(map[string]struct{}) + for _, item := range candidates { + kind := String(item, "type", "") + if kind != "function_call" && kind != "custom_tool_call" { + continue + } + callID := strings.TrimSpace(String(item, "call_id", "")) + if callID == "" { + continue + } + if _, ok := existingCalls[callID]; ok { + continue + } + if _, ok := existingOutputs[callID]; !ok { + continue + } + if _, ok := seen[callID]; ok { + continue + } + seen[callID] = struct{}{} + existingCalls[callID] = struct{}{} + filtered = append(filtered, clone(item)) + } + if len(filtered) == 0 { + return + } + + // Insert immediately before the first matching tool output. + insertAt := len(input) + for index, raw := range input { + item, ok := raw.(map[string]any) + if !ok { + continue + } + kind := String(item, "type", "") + if kind != "function_call_output" && kind != "custom_tool_call_output" { + continue + } + callID := strings.TrimSpace(String(item, "call_id", "")) + if callID == "" { + insertAt = index + break + } + if _, ok := seen[callID]; ok { + insertAt = index + break + } + // First output overall if none of the restored calls match yet. + if insertAt == len(input) { + insertAt = index + } + } + + missing := make([]any, 0, len(filtered)) + for _, item := range filtered { + missing = append(missing, item) + } + rewritten := make([]any, 0, len(input)+len(missing)) + rewritten = append(rewritten, input[:insertAt]...) + rewritten = append(rewritten, missing...) + rewritten = append(rewritten, input[insertAt:]...) + body["input"] = rewritten +} + +// pruneOrphanToolOutputs drops tool outputs whose call_id has no matching call. +func pruneOrphanToolOutputs(body map[string]any) { + input, ok := body["input"].([]any) + if !ok || len(input) == 0 { + return + } + calls := make(map[string]struct{}) + for _, raw := range input { + item, ok := raw.(map[string]any) + if !ok { + continue + } + switch String(item, "type", "") { + case "function_call", "custom_tool_call": + if callID := strings.TrimSpace(String(item, "call_id", "")); callID != "" { + calls[callID] = struct{}{} + } + } + } + var out []any + for index, raw := range input { + item, ok := raw.(map[string]any) + if !ok { + if out != nil { + out = append(out, raw) + } + continue + } + switch String(item, "type", "") { + case "function_call_output", "custom_tool_call_output": + callID := strings.TrimSpace(String(item, "call_id", "")) + _, matched := calls[callID] + if callID == "" || !matched { + if out == nil { + out = make([]any, 0, len(input)) + out = append(out, input[:index]...) + } + continue + } + } + if out != nil { + out = append(out, raw) + } + } + if out != nil { + body["input"] = out + } +} diff --git a/internal/openai/tool_replay_test.go b/internal/openai/tool_replay_test.go new file mode 100644 index 0000000..77442cf --- /dev/null +++ b/internal/openai/tool_replay_test.go @@ -0,0 +1,328 @@ +package openai + +import ( + "strings" + "testing" +) + +func TestToolReplayRestoresCallsFromPreviousResponseID(t *testing.T) { + cache := NewToolReplayCache(0, 0) + RememberCompletedResponse(cache, "grok-4.5", map[string]any{ + "id": "resp_turn1", + "model": "grok-4.5", + "output": []any{ + map[string]any{ + "id": "fc_1", "type": "function_call", "call_id": "call_1", + "name": "lookup", "arguments": `{"q":"weather"}`, + }, + }, + }, "") + + // Alma multi-turn: only tool output + previous_response_id. + wire, _, err := PrepareCompatibleResponsesWithCache(map[string]any{ + "model": "grok-4.5", + "previous_response_id": "resp_turn1", + "input": []any{ + map[string]any{"type": "function_call_output", "call_id": "call_1", "output": "sunny"}, + map[string]any{"type": "message", "role": "user", "content": []any{map[string]any{"type": "input_text", "text": "thanks"}}}, + }, + }, cache) + if err != nil { + t.Fatal(err) + } + if _, ok := wire["previous_response_id"]; ok { + t.Fatal("previous_response_id should be stripped before upstream") + } + input := wire["input"].([]any) + if len(input) < 2 { + t.Fatalf("input = %#v", input) + } + call := input[0].(map[string]any) + if call["type"] != "function_call" || call["call_id"] != "call_1" || call["name"] != "lookup" { + t.Fatalf("restored call = %#v", call) + } + output := input[1].(map[string]any) + if output["type"] != "function_call_output" || output["call_id"] != "call_1" { + t.Fatalf("tool output = %#v", output) + } +} + +func TestToolReplayExpandsItemReference(t *testing.T) { + cache := NewToolReplayCache(0, 0) + RememberCompletedResponse(cache, "grok-4.5", map[string]any{ + "id": "resp_turn1", + "model": "grok-4.5", + "output": []any{ + map[string]any{ + "id": "fc_item", "type": "function_call", "call_id": "call_ref", + "name": "search", "arguments": `{"q":"x"}`, + }, + }, + }, "") + + wire, _, err := PrepareCompatibleResponsesWithCache(map[string]any{ + "model": "grok-4.5", + "input": []any{ + map[string]any{"type": "message", "role": "user", "content": "continue"}, + map[string]any{"type": "item_reference", "id": "fc_item"}, + map[string]any{"type": "function_call_output", "call_id": "call_ref", "output": "ok"}, + }, + }, cache) + if err != nil { + t.Fatal(err) + } + input := wire["input"].([]any) + if len(input) != 3 { + t.Fatalf("input = %#v", input) + } + if String(input[1].(map[string]any), "type", "") != "function_call" { + t.Fatalf("expanded item = %#v", input[1]) + } +} + +func TestToolReplayPreservesNamespacedFunctionAlias(t *testing.T) { + cache := NewToolReplayCache(0, 0) + RememberCompletedResponse(cache, "grok-4.5", map[string]any{ + "id": "resp_ns", + "output": []any{map[string]any{ + "id": "fc_ns", "type": "function_call", "call_id": "call_ns", + "name": "fetch", "namespace": "mcp__github__", "arguments": `{}`, + }}, + }, "") + tools := []any{map[string]any{ + "type": "namespace", "name": "mcp__github__", + "tools": []any{map[string]any{"type": "function", "name": "fetch", "parameters": objectSchema()}}, + }} + tests := map[string]map[string]any{ + "previous_response_id": { + "model": "grok-4.5", "previous_response_id": "resp_ns", "tools": tools, + "input": []any{map[string]any{"type": "function_call_output", "call_id": "call_ns", "output": "ok"}}, + }, + "item_reference": { + "model": "grok-4.5", "tools": tools, + "input": []any{ + map[string]any{"type": "item_reference", "id": "fc_ns"}, + map[string]any{"type": "function_call_output", "call_id": "call_ns", "output": "ok"}, + }, + }, + } + for name, body := range tests { + t.Run(name, func(t *testing.T) { + wire, _, err := PrepareCompatibleResponsesWithCache(body, cache) + if err != nil { + t.Fatal(err) + } + input := wire["input"].([]any) + call := input[0].(map[string]any) + if got := String(call, "name", ""); got != "mcp__github__fetch" { + t.Fatalf("replayed call name = %q, want namespaced tool alias; call = %#v", got, call) + } + if _, ok := toolsByName(t, wire["tools"])["mcp__github__fetch"]; !ok { + t.Fatalf("namespaced tool missing from wire tools: %#v", wire["tools"]) + } + }) + } +} + +func TestNormalizeReplayItemsPreservesNamespace(t *testing.T) { + tests := []map[string]any{ + { + "type": "function_call", "call_id": "call_fn", "name": "fetch", + "namespace": "mcp__github__", "arguments": `{}`, + }, + { + "type": "custom_tool_call", "call_id": "call_custom", "name": "exec", + "namespace": "plugin__", "input": "echo ok", + }, + } + for _, item := range tests { + got := normalizeReplayItems([]map[string]any{item}) + if len(got) != 1 || got[0]["namespace"] != item["namespace"] { + t.Fatalf("normalized replay item lost namespace: got %#v, input %#v", got, item) + } + } +} + +func TestToolReplayPrunesOrphanOutputs(t *testing.T) { + cache := NewToolReplayCache(0, 0) + wire, _, err := PrepareCompatibleResponsesWithCache(map[string]any{ + "model": "grok-4.5", + "input": []any{ + map[string]any{"type": "item_reference", "id": "missing"}, + map[string]any{"type": "function_call_output", "call_id": "orphan", "output": "x"}, + map[string]any{"type": "message", "role": "user", "content": []any{map[string]any{"type": "input_text", "text": "hi"}}}, + }, + }, cache) + if err != nil { + t.Fatal(err) + } + input := wire["input"].([]any) + if len(input) != 1 || String(input[0].(map[string]any), "type", "") != "message" { + t.Fatalf("expected only message after orphan prune, got %#v", input) + } +} + +func TestPruneOrphanToolOutputsKeepsSliceWhenUnchanged(t *testing.T) { + input := []any{ + map[string]any{"type": "function_call", "call_id": "call_1", "name": "lookup", "arguments": `{}`}, + map[string]any{"type": "function_call_output", "call_id": "call_1", "output": "ok"}, + map[string]any{"type": "message", "role": "user", "content": "continue"}, + } + body := map[string]any{"input": input} + pruneOrphanToolOutputs(body) + got := body["input"].([]any) + if &got[0] != &input[0] { + t.Fatal("unchanged input was copied") + } +} + +func TestToolReplayRestoresFromStreamIndexedPreviousResponse(t *testing.T) { + cache := NewToolReplayCache(0, 0) + // Stream path: output_item.done indexes under prev-resp + item id. + RememberStreamToolCall(cache, "grok-4.5", map[string]any{ + "id": "fc_stream", "type": "function_call", "call_id": "call_stream", + "name": "lookup", "arguments": `{"q":"x"}`, + }, "resp_stream", "") + + // Alma multi-turn path: previous_response_id + item_reference + output. + wire, _, err := PrepareCompatibleResponsesWithCache(map[string]any{ + "model": "grok-4.5", + "previous_response_id": "resp_stream", + "input": []any{ + map[string]any{"type": "item_reference", "id": "fc_stream"}, + map[string]any{"type": "function_call_output", "call_id": "call_stream", "output": "result"}, + }, + }, cache) + if err != nil { + t.Fatal(err) + } + input := wire["input"].([]any) + if len(input) < 2 { + t.Fatalf("input = %#v", input) + } + call := input[0].(map[string]any) + if call["type"] != "function_call" || call["call_id"] != "call_stream" { + t.Fatalf("restored call = %#v", call) + } +} + +func TestToolReplayUsesClientModelNotUpstreamFreeAlias(t *testing.T) { + cache := NewToolReplayCache(0, 0) + // Upstream rewrites free-tier responses to grok-4.5-build-free. + RememberCompletedResponse(cache, "grok-4.5", map[string]any{ + "id": "resp_free", + "model": "grok-4.5-build-free", + "output": []any{ + map[string]any{ + "id": "fc_free", "type": "function_call", "call_id": "call_free", + "name": "lookup", "arguments": `{"q":"x"}`, + }, + }, + }, "") + + // Next Alma turn still asks for the client model name. + wire, _, err := PrepareCompatibleResponsesWithCache(map[string]any{ + "model": "grok-4.5", + "previous_response_id": "resp_free", + "input": []any{ + map[string]any{"type": "function_call_output", "call_id": "call_free", "output": "ok"}, + }, + }, cache) + if err != nil { + t.Fatal(err) + } + input := wire["input"].([]any) + if len(input) < 1 || String(input[0].(map[string]any), "type", "") != "function_call" { + t.Fatalf("client model lookup missed free-tier cache entry: %#v", input) + } +} + +func TestPrepareCompatibleResponsesHardensArgumentAndOutputShapes(t *testing.T) { + wire, _, err := PrepareCompatibleResponses(map[string]any{ + "model": "grok-4.5", + "input": []any{ + map[string]any{ + "type": "function_call", "call_id": "c1", "name": "lookup", + "arguments": map[string]any{"q": "x"}, + }, + map[string]any{ + "type": "function_call_output", "call_id": "c1", + "output": []any{ + map[string]any{"type": "output_text", "text": "line1"}, + map[string]any{"type": "output_text", "text": "line2"}, + }, + }, + }, + "tools": []any{map[string]any{"type": "function", "name": "lookup", "parameters": objectSchema()}}, + }) + if err != nil { + t.Fatal(err) + } + input := wire["input"].([]any) + call := input[0].(map[string]any) + args, ok := call["arguments"].(string) + if !ok || !strings.Contains(args, `"q"`) { + t.Fatalf("arguments not stringified: %#v", call["arguments"]) + } + output := input[1].(map[string]any) + if got, ok := output["output"].(string); !ok || !strings.Contains(got, "line1") || !strings.Contains(got, "line2") { + t.Fatalf("output not flattened: %#v", output["output"]) + } +} + +func TestPrepareCompatibleResponsesDropsServerToolHistory(t *testing.T) { + wire, _, err := PrepareCompatibleResponses(map[string]any{ + "model": "grok-4.5", + "input": []any{ + map[string]any{"type": "message", "role": "user", "content": []any{map[string]any{"type": "input_text", "text": "hi"}}}, + map[string]any{"type": "web_search_call", "id": "ws_1", "status": "completed"}, + map[string]any{"type": "x_search_call", "id": "xs_1", "status": "completed"}, + map[string]any{"type": "message", "role": "user", "content": []any{map[string]any{"type": "input_text", "text": "again"}}}, + }, + }) + if err != nil { + t.Fatal(err) + } + input := wire["input"].([]any) + if len(input) != 2 { + t.Fatalf("expected server tool history dropped, got %#v", input) + } + for _, raw := range input { + if String(raw.(map[string]any), "type", "") != "message" { + t.Fatalf("unexpected item %#v", raw) + } + } +} + +func TestPrepareCompatibleResponsesStripsNullReasoningFields(t *testing.T) { + // Live 422: reasoning with content:null fails ModelInput untagged enum. + wire, _, err := PrepareCompatibleResponses(map[string]any{ + "model": "grok-4.5", + "input": []any{ + map[string]any{ + "type": "reasoning", "id": "rs_1", + "summary": []any{map[string]any{"type": "summary_text", "text": "think"}}, + "content": nil, + "encrypted_content": nil, + }, + map[string]any{"type": "message", "role": "user", "content": []any{map[string]any{"type": "input_text", "text": "hi"}}}, + }, + }) + if err != nil { + t.Fatal(err) + } + input := wire["input"].([]any) + if len(input) != 2 { + t.Fatalf("input = %#v", input) + } + reasoning := input[0].(map[string]any) + if reasoning["type"] != "reasoning" { + t.Fatalf("reasoning dropped unexpectedly: %#v", reasoning) + } + if _, ok := reasoning["content"]; ok { + t.Fatalf("content null must be stripped: %#v", reasoning) + } + if _, ok := reasoning["encrypted_content"]; ok { + t.Fatalf("encrypted_content must be stripped: %#v", reasoning) + } +} diff --git a/internal/server/server.go b/internal/server/server.go index ab67bb0..c9e5b58 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -341,6 +341,7 @@ func (s *Server) responses(w http.ResponseWriter, r *http.Request) { model := openai.String(body, "model", "") affinity := requestAffinity(r, body) convID := conversationID(affinity) + promptCacheKey := openai.String(body, "prompt_cache_key", "") if !openai.IsStreaming(body) { payload, err := s.client.DoJSON(r.Context(), http.MethodPost, "responses", wire, affinity, convID, fmt.Sprint(wire["model"]), true) if err != nil { @@ -350,11 +351,15 @@ func (s *Server) responses(w http.ResponseWriter, r *http.Request) { if native { writeJSON(w, http.StatusOK, payload) } else { - writeJSON(w, http.StatusOK, compat.NormalizeResponse(payload, model)) + normalized := compat.NormalizeResponse(payload, model) + // Index tool calls so Alma can continue with previous_response_id + // / item_reference without re-sending the matching function_call. + openai.RememberCompletedResponse(openai.DefaultToolReplay, model, normalized, promptCacheKey) + writeJSON(w, http.StatusOK, normalized) } return } - s.streamResponses(w, r, wire, affinity, convID, model, native, compat) + s.streamResponses(w, r, wire, affinity, convID, model, native, compat, promptCacheKey) } func (s *Server) messages(w http.ResponseWriter, r *http.Request) { @@ -439,7 +444,7 @@ func (s *Server) streamChat(w http.ResponseWriter, r *http.Request, wire map[str flush() } -func (s *Server) streamResponses(w http.ResponseWriter, r *http.Request, wire map[string]any, affinity auth.Affinity, convID, model string, native bool, compat *openai.ResponsesCompatibility) { +func (s *Server) streamResponses(w http.ResponseWriter, r *http.Request, wire map[string]any, affinity auth.Affinity, convID, model string, native bool, compat *openai.ResponsesCompatibility, promptCacheKey string) { stream, err := s.client.OpenStream(r.Context(), "responses", wire, affinity, convID, fmt.Sprint(wire["model"]), true) if err != nil { s.writeClientError(w, err) @@ -450,6 +455,10 @@ func (s *Server) streamResponses(w http.ResponseWriter, r *http.Request, wire ma flush := flusher(w) timing := grok.RequestTimingFromContext(r.Context()) timing.MarkDownstreamFlush(false) + // Accumulate tool calls from output_item.done, then index under + // prev-resp:{response.id} when response.completed arrives (done events + // often omit response_id, so early prev-resp writes can miss). + replay := &streamToolReplayState{model: model, promptCacheKey: promptCacheKey} for { event, ok, nextErr := stream.Next() if nextErr != nil { @@ -487,6 +496,7 @@ func (s *Server) streamResponses(w http.ResponseWriter, r *http.Request, wire ma translatedEvent.Event = openai.EventType("", data) } } + replay.handle(translatedEvent.Event, translatedEvent.Data) if err := writeRawSSE(w, translatedEvent); err != nil { return } @@ -504,6 +514,78 @@ func (s *Server) streamResponses(w http.ResponseWriter, r *http.Request, wire ma } } +// streamToolReplayState accumulates tool calls from a single Responses stream +// (collect on output_item.done, cache on response.completed). +type streamToolReplayState struct { + model string + promptCacheKey string + responseID string + // items holds function/custom tool calls from output_item.done, ordered. + items []map[string]any +} + +func (s *streamToolReplayState) handle(event string, data []byte) { + if s == nil { + return + } + var payload map[string]any + if err := json.Unmarshal(data, &payload); err != nil { + return + } + kind := openai.EventType(event, payload) + if id := openai.String(payload, "response_id", ""); id != "" { + s.responseID = id + } + if response, ok := payload["response"].(map[string]any); ok { + if id := openai.String(response, "id", ""); id != "" { + s.responseID = id + } + } + switch kind { + case "response.created", "response.in_progress": + // Learn response id only. + case "response.output_item.done": + // Accumulate tool calls for patching completed.output if empty. + item, _ := payload["item"].(map[string]any) + if item == nil { + return + } + switch openai.String(item, "type", "") { + case "function_call", "custom_tool_call": + s.items = append(s.items, item) + // Index item:{id} immediately so item_reference works even when + // Alma continues before we see response.completed (partial turns). + // Do NOT write prev-resp here: done events often omit response_id + // and the session key is only committed on completed. + if id := openai.String(item, "id", ""); id != "" { + openai.RememberStreamToolCall(openai.DefaultToolReplay, s.model, item, "", "") + } + } + case "response.completed": + // Patch empty completed.output from collected done items, then index + // under prev-resp / item / session keys. + response, _ := payload["response"].(map[string]any) + if response == nil { + response = map[string]any{} + } + if id := openai.String(response, "id", ""); id != "" { + s.responseID = id + } else if s.responseID != "" { + response["id"] = s.responseID + } + if output, ok := response["output"].([]any); !ok || len(output) == 0 { + if len(s.items) > 0 { + patched := make([]any, 0, len(s.items)) + for _, item := range s.items { + patched = append(patched, item) + } + response["output"] = patched + } + } + openai.RememberCompletedResponse(openai.DefaultToolReplay, s.model, response, s.promptCacheKey) + } +} + func (s *Server) streamAnthropic(w http.ResponseWriter, r *http.Request, wire map[string]any, affinity auth.Affinity, convID, model string, options anthropic.ResponseOptions) { stream, err := s.client.OpenStream(r.Context(), "responses", wire, affinity, convID, fmt.Sprint(wire["model"]), true) if err != nil {