From 242f66325cac4d5cc392dcc001c74ef7215426c0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Wed, 21 Oct 2026 18:43:36 +0200 Subject: [PATCH 01/11] test(anthropic): script server fallback streams MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add replies for a fallback before any output, a fallback mid-output after a client tool call, and sticky routing, checked against the SDK's accumulator. Signed-off-by: Arnaud Héritier --- .../anthropic/anthropictest/server.go | 84 +++++++++++++++++-- .../anthropic/anthropictest/server_test.go | 42 ++++++++++ 2 files changed, 121 insertions(+), 5 deletions(-) diff --git a/pkg/model/provider/anthropic/anthropictest/server.go b/pkg/model/provider/anthropic/anthropictest/server.go index 8b58f72c02..49e0336f45 100644 --- a/pkg/model/provider/anthropic/anthropictest/server.go +++ b/pkg/model/provider/anthropic/anthropictest/server.go @@ -70,7 +70,7 @@ func (s *Server) Requests() []Request { // TextReply is a plain text answer reporting inputTokens. func TextReply(id, text string, inputTokens int64) string { return body( - messageStart(id, inputTokens), + messageStart(id, "claude-test", inputTokens), event("content_block_start", map[string]any{"index": 0, "content_block": map[string]any{"type": "text", "text": ""}}), event("content_block_delta", map[string]any{"index": 0, "delta": map[string]any{"type": "text_delta", "text": text}}), event("content_block_stop", map[string]any{"index": 0}), @@ -82,7 +82,7 @@ func TextReply(id, text string, inputTokens int64) string { // ThinkingTextReply is signed thinking followed by a text answer. func ThinkingTextReply(id, thinking, signature, text string, inputTokens int64) string { return body( - messageStart(id, inputTokens), + messageStart(id, "claude-test", inputTokens), event("content_block_start", map[string]any{"index": 0, "content_block": map[string]any{"type": "thinking", "thinking": "", "signature": ""}}), event("content_block_delta", map[string]any{"index": 0, "delta": map[string]any{"type": "thinking_delta", "thinking": thinking}}), event("content_block_delta", map[string]any{"index": 0, "delta": map[string]any{"type": "signature_delta", "signature": signature}}), @@ -103,7 +103,7 @@ func CompactionReply(id, summary, signature string, iterations ...[2]int64) stri usage = append(usage, map[string]any{"type": "compaction", "input_tokens": it[0], "output_tokens": it[1]}) } return body( - messageStart(id, 0), + messageStart(id, "claude-test", 0), event("content_block_start", map[string]any{"index": 0, "content_block": map[string]any{"type": "compaction", "content": nil, "encrypted_content": "", "signature": nil}}), event("content_block_delta", map[string]any{"index": 0, "delta": map[string]any{"type": "compaction_delta", "content": summary, "encrypted_content": "enc-" + id}}), event("content_block_delta", map[string]any{"index": 0, "delta": map[string]any{"type": "signature_delta", "signature": signature}}), @@ -113,6 +113,80 @@ func CompactionReply(id, summary, signature string, iterations ...[2]int64) stri ) } +// FallbackBeforeOutputReply is a turn that from declined before any output: +// message_start already names to, and the fallback marker comes first. +func FallbackBeforeOutputReply(id, from, to, text string) string { + events := append([]sse{messageStart(id, to, 0)}, fallbackBlock(0, from, to)...) + events = append(events, textBlock(1, text)...) + events = append(events, messageDelta("end_turn", fallbackUsage(from, to)), event("message_stop", nil)) + return body(events...) +} + +// FallbackMidOutputReply is a turn that from declined mid-output, after its +// signed thinking, the text before and a complete client tool call toolID; +// to continues with the text after. +func FallbackMidOutputReply(id, from, to, before, after, toolID string) string { + events := []sse{ + messageStart(id, from, 0), + event("content_block_start", map[string]any{"index": 0, "content_block": map[string]any{"type": "thinking", "thinking": "", "signature": ""}}), + event("content_block_delta", map[string]any{"index": 0, "delta": map[string]any{"type": "thinking_delta", "thinking": "declined plan"}}), + event("content_block_delta", map[string]any{"index": 0, "delta": map[string]any{"type": "signature_delta", "signature": "sig-declined"}}), + event("content_block_stop", map[string]any{"index": 0}), + } + events = append(events, textBlock(1, before)...) + events = append(events, + event("content_block_start", map[string]any{"index": 2, "content_block": map[string]any{"type": "tool_use", "id": toolID, "name": "record", "input": map[string]any{}}}), + event("content_block_delta", map[string]any{"index": 2, "delta": map[string]any{"type": "input_json_delta", "partial_json": "{}"}}), + event("content_block_stop", map[string]any{"index": 2}), + ) + events = append(events, fallbackBlock(3, from, to)...) + events = append(events, textBlock(4, after)...) + events = append(events, messageDelta("end_turn", fallbackUsage(from, to)), event("message_stop", nil)) + return body(events...) +} + +// StickyReply is a turn sticky routing sent straight to the fallback model: +// no marker, only the model and a fallback_message iteration show it. +func StickyReply(id, to, text string) string { + events := []sse{messageStart(id, to, 0)} + events = append(events, textBlock(0, text)...) + events = append(events, messageDelta("end_turn", map[string]any{ + "input_tokens": 0, "output_tokens": 7, + "iterations": []any{iteration("fallback_message", to)}, + }), event("message_stop", nil)) + return body(events...) +} + +// fallbackBlock is a fallback marker: a block start and stop with no deltas. +func fallbackBlock(index int, from, to string) []sse { + return []sse{ + event("content_block_start", map[string]any{"index": index, "content_block": map[string]any{ + "type": "fallback", "from": map[string]any{"model": from}, "to": map[string]any{"model": to}, + "trigger": map[string]any{"type": "refusal", "category": "cyber"}, + }}), + event("content_block_stop", map[string]any{"index": index}), + } +} + +func textBlock(index int, text string) []sse { + return []sse{ + event("content_block_start", map[string]any{"index": index, "content_block": map[string]any{"type": "text", "text": ""}}), + event("content_block_delta", map[string]any{"index": index, "delta": map[string]any{"type": "text_delta", "text": text}}), + event("content_block_stop", map[string]any{"index": index}), + } +} + +func fallbackUsage(from, to string) map[string]any { + return map[string]any{ + "input_tokens": 0, "output_tokens": 7, + "iterations": []any{iteration("message", from), iteration("fallback_message", to)}, + } +} + +func iteration(typ, model string) map[string]any { + return map[string]any{"type": typ, "model": model, "input_tokens": 10, "output_tokens": 7, "cache_read_input_tokens": 0, "cache_creation_input_tokens": 0} +} + type sse struct { name string data map[string]any @@ -126,9 +200,9 @@ func event(name string, data map[string]any) sse { return sse{name, data} } -func messageStart(id string, inputTokens int64) sse { +func messageStart(id, model string, inputTokens int64) sse { return event("message_start", map[string]any{"message": map[string]any{ - "id": id, "type": "message", "role": "assistant", "model": "claude-test", "content": []any{}, + "id": id, "type": "message", "role": "assistant", "model": model, "content": []any{}, "usage": map[string]any{"input_tokens": inputTokens, "output_tokens": 0}, }}) } diff --git a/pkg/model/provider/anthropic/anthropictest/server_test.go b/pkg/model/provider/anthropic/anthropictest/server_test.go index 406d800ac3..1260b75b3c 100644 --- a/pkg/model/provider/anthropic/anthropictest/server_test.go +++ b/pkg/model/provider/anthropic/anthropictest/server_test.go @@ -6,6 +6,8 @@ import ( "strings" "testing" + "github.com/anthropics/anthropic-sdk-go" + "github.com/anthropics/anthropic-sdk-go/option" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -32,3 +34,43 @@ func TestServerRepliesAndRecords(t *testing.T) { assert.Equal(t, []string{"a", "b"}, got[0].Betas) assert.Equal(t, "m", got[0].Body["model"]) } + +// The fallback replies must parse with the SDK, which relabels the message +// with the model that served it. +func TestFallbackRepliesAccumulate(t *testing.T) { + t.Parallel() + + for name, tc := range map[string]struct { + reply string + markers int + }{ + "before output": {reply: anthropictest.FallbackBeforeOutputReply("m1", "claude-a", "claude-b", "hi"), markers: 1}, + "mid output": {reply: anthropictest.FallbackMidOutputReply("m2", "claude-a", "claude-b", "before", "after", "t1"), markers: 1}, + "sticky": {reply: anthropictest.StickyReply("m3", "claude-b", "hi")}, + } { + t.Run(name, func(t *testing.T) { + t.Parallel() + srv := anthropictest.NewServer(t, tc.reply) + client := anthropic.NewClient(option.WithBaseURL(srv.URL()), option.WithAPIKey("k")) + stream := client.Beta.Messages.NewStreaming(t.Context(), anthropic.BetaMessageNewParams{ + Model: "claude-a", MaxTokens: 16, + Messages: []anthropic.BetaMessageParam{anthropic.NewBetaUserMessage(anthropic.NewBetaTextBlock("q"))}, + }) + var msg anthropic.BetaMessage + for stream.Next() { + require.NoError(t, msg.Accumulate(stream.Current())) + } + require.NoError(t, stream.Err()) + + assert.Equal(t, anthropic.Model("claude-b"), msg.Model) + markers := 0 + for _, block := range msg.Content { + if block.Type == "fallback" { + markers++ + } + } + assert.Equal(t, tc.markers, markers) + assert.NotEmpty(t, msg.Usage.Iterations) + }) + } +} From fa8661c38ae4ef1ba4cede31fb83f3023d8daabe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Wed, 21 Oct 2026 18:48:34 +0200 Subject: [PATCH 02/11] feat(anthropic): report server fallback markers and the serving model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Each fallback content block reaches the runtime as a ServerFallback marker. The stream reports the model a server fallback handed the turn to, from the marker or, for sticky routing, from a fallback_message iteration, through a new ServedStream interface read once the stream ends. Wrappers forward it instead of snapshotting it, and a normal turn reports nothing, so aliases and pricing lookups don't change. Signed-off-by: Arnaud Héritier --- pkg/chat/chat.go | 10 ++ .../anthropic/anthropictest/server.go | 21 ++++ pkg/model/provider/anthropic/beta_adapter.go | 31 ++++++ .../anthropic/server_fallback_test.go | 103 ++++++++++++++++++ pkg/model/provider/contracts/contracts.go | 8 ++ pkg/model/provider/instrument.go | 17 ++- pkg/model/provider/provider.go | 13 +++ pkg/model/provider/rulebased/client.go | 7 ++ pkg/model/provider/serving_test.go | 52 +++++++++ pkg/runtime/fallback.go | 3 + 10 files changed, 264 insertions(+), 1 deletion(-) create mode 100644 pkg/model/provider/anthropic/server_fallback_test.go diff --git a/pkg/chat/chat.go b/pkg/chat/chat.go index e2a49677e1..3e13027376 100644 --- a/pkg/chat/chat.go +++ b/pkg/chat/chat.go @@ -198,6 +198,16 @@ type MessageDelta struct { // pointer — a scalar field silently dropped every blob but the last. See // [MediaDelta]. Media []MediaDelta `json:"media,omitempty"` + // ServerFallback marks the point where the provider handed the turn to + // a fallback model: work streamed before it is superseded. + ServerFallback *ServerFallback `json:"server_fallback,omitempty"` +} + +// ServerFallback is a provider-side handoff from one model to the next +// within a single response. +type ServerFallback struct { + From string `json:"from"` + To string `json:"to"` } // MessageStreamChoice represents a choice in a streaming response diff --git a/pkg/model/provider/anthropic/anthropictest/server.go b/pkg/model/provider/anthropic/anthropictest/server.go index 49e0336f45..836b942137 100644 --- a/pkg/model/provider/anthropic/anthropictest/server.go +++ b/pkg/model/provider/anthropic/anthropictest/server.go @@ -145,6 +145,27 @@ func FallbackMidOutputReply(id, from, to, before, after, toolID string) string { return body(events...) } +// FallbackChainReply is a turn handed down a chain of models: each one +// writes a short text, then declines to the next, which ends the turn. +func FallbackChainReply(id string, models ...string) string { + events := []sse{messageStart(id, models[0], 0)} + var iterations []any + index := 0 + for i, model := range models { + events = append(events, textBlock(index, "from "+model)...) + index++ + if i == len(models)-1 { + iterations = append(iterations, iteration("fallback_message", model)) + break + } + iterations = append(iterations, iteration("message", model)) + events = append(events, fallbackBlock(index, model, models[i+1])...) + index++ + } + events = append(events, messageDelta("end_turn", map[string]any{"input_tokens": 0, "output_tokens": 7, "iterations": iterations}), event("message_stop", nil)) + return body(events...) +} + // StickyReply is a turn sticky routing sent straight to the fallback model: // no marker, only the model and a fallback_message iteration show it. func StickyReply(id, to, text string) string { diff --git a/pkg/model/provider/anthropic/beta_adapter.go b/pkg/model/provider/anthropic/beta_adapter.go index 8a0d9bbec8..cfa171789e 100644 --- a/pkg/model/provider/anthropic/beta_adapter.go +++ b/pkg/model/provider/anthropic/beta_adapter.go @@ -9,6 +9,8 @@ import ( "github.com/anthropics/anthropic-sdk-go/packages/ssestream" "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/model/provider/base" + "github.com/docker/docker-agent/pkg/model/provider/contracts" "github.com/docker/docker-agent/pkg/tools" ) @@ -27,6 +29,10 @@ type betaStreamAdapter struct { message anthropic.BetaMessage rawLost bool requestContext json.RawMessage + // config is the dispatching client's; served is the model a server + // fallback handed the turn to, empty when the dispatched model served it. + config base.Config + served string } // newBetaStreamAdapter creates a new Beta stream adapter @@ -35,6 +41,7 @@ func (c *Client) newBetaStreamAdapter(stream *ssestream.Stream[anthropic.BetaRaw retryableStream: retryableStream[anthropic.BetaRawMessageStreamEventUnion]{stream: stream}, trackUsage: trackUsage, toolIDByBlock: map[int64]string{}, + config: c.Config, } } @@ -77,6 +84,9 @@ func (a *betaStreamAdapter) Recv() (chat.MessageStreamResponse, error) { }, } response.Choices[0].Delta.ToolCalls = []tools.ToolCall{toolCall} + case anthropic.BetaFallbackBlock: + a.served = block.To.Model + response.Choices[0].Delta.ServerFallback = &chat.ServerFallback{From: block.From.Model, To: block.To.Model} case anthropic.BetaThinkingBlock: if block.Thinking != "" { response.Choices[0].Delta.ReasoningContent = block.Thinking @@ -110,6 +120,13 @@ func (a *betaStreamAdapter) Recv() (chat.MessageStreamResponse, error) { } case anthropic.BetaRawMessageDeltaEvent: a.stopReason = eventVariant.Delta.StopReason + // Sticky routing sends no marker: only a fallback_message iteration + // shows that a fallback model served the turn. + for _, it := range a.message.Usage.Iterations { + if it.Type == "fallback_message" && a.served == "" { + a.served = it.Model + } + } if a.trackUsage { response.Usage = betaUsageFromMessage(a.message.Usage) } @@ -146,6 +163,20 @@ func (a *betaStreamAdapter) accumulate(event anthropic.BetaRawMessageStreamEvent } } +// ServedBy reports the model a server fallback handed the turn to. +func (a *betaStreamAdapter) ServedBy() (contracts.ServingModel, bool) { + if a.served == "" { + return contracts.ServingModel{}, false + } + cfg := a.config + cfg.ModelConfig.Model = a.served + // The primary's display name and cost overrides don't describe the + // fallback model, which is priced from the catalogue. + cfg.ModelConfig.DisplayModel = "" + cfg.ModelConfig.Cost = nil + return contracts.ServingModel{ID: cfg.ID(), Config: cfg}, true +} + // Close closes the Beta stream func (a *betaStreamAdapter) Close() { a.close() diff --git a/pkg/model/provider/anthropic/server_fallback_test.go b/pkg/model/provider/anthropic/server_fallback_test.go new file mode 100644 index 0000000000..d481a6353f --- /dev/null +++ b/pkg/model/provider/anthropic/server_fallback_test.go @@ -0,0 +1,103 @@ +package anthropic + +import ( + "errors" + "io" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/config/latest" + "github.com/docker/docker-agent/pkg/model/provider/anthropic/anthropictest" + "github.com/docker/docker-agent/pkg/model/provider/contracts" +) + +func fallbackClient(srv *scriptedServer) *Client { + return newScriptedClient(srv, latest.ModelConfig{ + Provider: "anthropic", Model: "claude-opus-4-8", + ProviderOpts: map[string]any{"fallbacks": []any{"claude-sonnet-4-6"}}, + }) +} + +// drainMarkers reads a stream to its end and returns its server fallback +// markers. +func drainMarkers(t *testing.T, stream chat.MessageStream) []chat.ServerFallback { + t.Helper() + var markers []chat.ServerFallback + for { + resp, err := stream.Recv() + if errors.Is(err, io.EOF) { + return markers + } + require.NoError(t, err) + for _, c := range resp.Choices { + if c.Delta.ServerFallback != nil { + markers = append(markers, *c.Delta.ServerFallback) + } + } + } +} + +func TestServerFallbackStream(t *testing.T) { + t.Parallel() + + const from, to = "claude-opus-4-8", "claude-sonnet-4-6" + for name, tc := range map[string]struct { + reply string + markers []chat.ServerFallback + served string + }{ + "before output": { + reply: anthropictest.FallbackBeforeOutputReply("m", from, to, "hi"), + markers: []chat.ServerFallback{{From: from, To: to}}, served: to, + }, + "mid output": { + reply: anthropictest.FallbackMidOutputReply("m", from, to, "before", "after", "t1"), + markers: []chat.ServerFallback{{From: from, To: to}}, served: to, + }, + "two hops": { + reply: anthropictest.FallbackChainReply("m", from, "claude-opus-4-7", to), + markers: []chat.ServerFallback{{From: from, To: "claude-opus-4-7"}, {From: "claude-opus-4-7", To: to}}, served: to, + }, + "sticky routing": {reply: anthropictest.StickyReply("m", to, "hi"), served: to}, + // message_start names "claude-test", not the dispatched model: a + // response model alone is no fallback. + "no fallback": {reply: anthropictest.TextReply("m", "hi", 10)}, + } { + t.Run(name, func(t *testing.T) { + t.Parallel() + stream, err := fallbackClient(newScriptedServer(t, tc.reply)).CreateChatCompletionStream(t.Context(), []chat.Message{user("q")}, nil) + require.NoError(t, err) + defer stream.Close() + + assert.Equal(t, tc.markers, drainMarkers(t, stream)) + served, ok := stream.(contracts.ServedStream) + require.True(t, ok) + model, ok := served.ServedBy() + if tc.served == "" { + assert.False(t, ok) + return + } + require.True(t, ok) + assert.Equal(t, "anthropic/"+tc.served, model.ID.String()) + assert.Equal(t, tc.served, model.Config.ModelConfig.Model) + }) + } + + t.Run("stream cut after the marker", func(t *testing.T) { + t.Parallel() + reply := anthropictest.FallbackMidOutputReply("m", from, to, "before", "after", "t1") + marker := strings.Index(reply, `"type":"fallback"`) + reply = reply[:marker+strings.Index(reply[marker:], "event: content_block_start")] + stream, err := fallbackClient(newScriptedServer(t, reply)).CreateChatCompletionStream(t.Context(), []chat.Message{user("q")}, nil) + require.NoError(t, err) + defer stream.Close() + assert.Equal(t, []chat.ServerFallback{{From: from, To: to}}, drainMarkers(t, stream)) + model, ok := stream.(contracts.ServedStream).ServedBy() + require.True(t, ok, "the handoff happened before the cut") + assert.Equal(t, to, model.Config.ModelConfig.Model) + }) +} diff --git a/pkg/model/provider/contracts/contracts.go b/pkg/model/provider/contracts/contracts.go index 89287a00bc..0cb4a8d36b 100644 --- a/pkg/model/provider/contracts/contracts.go +++ b/pkg/model/provider/contracts/contracts.go @@ -154,6 +154,14 @@ type ServingStream interface { ServingModel() ServingModel } +// ServedStream reports, once the stream has ended, the model that served the +// response when it isn't the dispatched one (e.g. a server-side fallback). +// Wrappers must forward it rather than snapshot it at creation. +type ServedStream interface { + chat.MessageStream + ServedBy() (ServingModel, bool) +} + // EmbeddingProvider is a provider that supports embeddings. type EmbeddingProvider interface { Provider diff --git a/pkg/model/provider/instrument.go b/pkg/model/provider/instrument.go index 7d8712a2e5..6b980b7dfb 100644 --- a/pkg/model/provider/instrument.go +++ b/pkg/model/provider/instrument.go @@ -135,7 +135,10 @@ func (t *tracedChat) CreateChatCompletionStream(ctx context.Context, messages [] } wrapped := genai.WrapStream(span, stream) if serving, ok := stream.(ServingStream); ok { - return &tracedServingStream{MessageStream: wrapped, serving: serving.ServingModel()}, nil + return &tracedServingStream{MessageStream: wrapped, serving: serving.ServingModel(), inner: stream}, nil + } + if _, ok := stream.(ServedStream); ok { + return &tracedServedStream{MessageStream: wrapped, inner: stream}, nil } return wrapped, nil } @@ -293,8 +296,20 @@ type tracedServingStream struct { chat.MessageStream serving ServingModel + inner chat.MessageStream } func (s *tracedServingStream) ServingModel() ServingModel { return s.serving } + +func (s *tracedServingStream) ServedBy() (ServingModel, bool) { return ServedBy(s.inner) } + +// tracedServedStream forwards ServedBy, which is only final once the stream ends. +type tracedServedStream struct { + chat.MessageStream + + inner chat.MessageStream +} + +func (s *tracedServedStream) ServedBy() (ServingModel, bool) { return ServedBy(s.inner) } diff --git a/pkg/model/provider/provider.go b/pkg/model/provider/provider.go index a9b47ed9e8..aab1c35276 100644 --- a/pkg/model/provider/provider.go +++ b/pkg/model/provider/provider.go @@ -11,6 +11,7 @@ package provider import ( "context" + "github.com/docker/docker-agent/pkg/chat" "github.com/docker/docker-agent/pkg/config/latest" "github.com/docker/docker-agent/pkg/environment" "github.com/docker/docker-agent/pkg/model/provider/contracts" @@ -26,6 +27,18 @@ type ServingModel = contracts.ServingModel // ServingStream reports the serving endpoint when it differs from the dispatched provider. type ServingStream = contracts.ServingStream +// ServedStream reports a model that served the response in place of the dispatched one. +type ServedStream = contracts.ServedStream + +// ServedBy reports the model that served a finished stream in place of the +// dispatched one; false when the stream doesn't say or nothing changed. +func ServedBy(stream chat.MessageStream) (ServingModel, bool) { + if served, ok := stream.(ServedStream); ok { + return served.ServedBy() + } + return ServingModel{}, false +} + // EmbeddingProvider is a provider that supports embeddings. type EmbeddingProvider = contracts.EmbeddingProvider diff --git a/pkg/model/provider/rulebased/client.go b/pkg/model/provider/rulebased/client.go index 421d6e5123..c8bebcc92b 100644 --- a/pkg/model/provider/rulebased/client.go +++ b/pkg/model/provider/rulebased/client.go @@ -228,3 +228,10 @@ type routedStream struct { func (s *routedStream) ServingModel() contracts.ServingModel { return s.serving } + +func (s *routedStream) ServedBy() (contracts.ServingModel, bool) { + if served, ok := s.MessageStream.(contracts.ServedStream); ok { + return served.ServedBy() + } + return contracts.ServingModel{}, false +} diff --git a/pkg/model/provider/serving_test.go b/pkg/model/provider/serving_test.go index 6a91e077dd..13d2be41bd 100644 --- a/pkg/model/provider/serving_test.go +++ b/pkg/model/provider/serving_test.go @@ -96,3 +96,55 @@ func TestServingModelThroughRegistry(t *testing.T) { }) } } + +// servedTestStream learns the model that served it while it streams. +type servedTestStream struct { + instrumentTestStream + + served string +} + +func (s *servedTestStream) Recv() (chat.MessageStreamResponse, error) { + s.served = "fallback" + return s.instrumentTestStream.Recv() +} + +func (s *servedTestStream) ServedBy() (ServingModel, bool) { + if s.served == "" { + return ServingModel{}, false + } + return ServingModel{ID: modelsdev.NewID("test", s.served)}, true +} + +// ServedBy is final only once the stream has ended, so wrappers must read it +// late rather than snapshot it at creation. +func TestServedByThroughRegistry(t *testing.T) { + t.Parallel() + + for name, routed := range map[string]bool{"direct": false, "routed": true} { + t.Run(name, func(t *testing.T) { + t.Parallel() + underlying := &servedTestStream{} + registry := NewRegistry(map[string]Factory{"test": func(_ context.Context, cfg *latest.ModelConfig, _ environment.Provider, _ ...options.Opt) (Provider, error) { + return ®istryServingProvider{Config: base.Config{ModelConfig: *cfg}, stream: underlying}, nil + }}) + cfg := &latest.ModelConfig{Provider: "test", Model: "child"} + if routed { + cfg.Model = "router" + cfg.Routing = []latest.RoutingRule{{Model: "test/child", Examples: []string{"poetry"}}} + } + p, err := registry.New(t.Context(), cfg, nil) + require.NoError(t, err) + stream, err := p.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "poetry"}}, nil) + require.NoError(t, err) + + _, ok := ServedBy(stream) + assert.False(t, ok, "nothing served yet") + _, _ = stream.Recv() + served, ok := ServedBy(stream) + require.True(t, ok) + assert.Equal(t, modelsdev.NewID("test", "fallback"), served.ID) + stream.Close() + }) + } +} diff --git a/pkg/runtime/fallback.go b/pkg/runtime/fallback.go index 66ba8d6a83..cff0e7f5ef 100644 --- a/pkg/runtime/fallback.go +++ b/pkg/runtime/fallback.go @@ -580,6 +580,9 @@ func finalModelFor(serving *provider.ServingModel, dispatched provider.Provider) } func servingModelFrom(stream chat.MessageStream) *provider.ServingModel { + if served, ok := provider.ServedBy(stream); ok { + return &served + } if stream, ok := stream.(provider.ServingStream); ok { return new(stream.ServingModel()) } From 885a9ace5fc8e8ddc742b25f2b5c7ea5aea6c128 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Wed, 21 Oct 2026 18:50:31 +0200 Subject: [PATCH 03/11] feat(runtime): drop work superseded by a server fallback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit On a server fallback marker, the tool calls and thinking streamed so far belong to the model that declined: drop them so they are never executed or replayed, keep the text, and warn which model took over. Signed-off-by: Arnaud Héritier --- pkg/runtime/server_fallback_test.go | 84 +++++++++++++++++++++++++++++ pkg/runtime/streaming.go | 17 ++++++ 2 files changed, 101 insertions(+) create mode 100644 pkg/runtime/server_fallback_test.go diff --git a/pkg/runtime/server_fallback_test.go b/pkg/runtime/server_fallback_test.go new file mode 100644 index 0000000000..42181b390e --- /dev/null +++ b/pkg/runtime/server_fallback_test.go @@ -0,0 +1,84 @@ +package runtime + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/agent" + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/session" +) + +func (b *streamBuilder) AddServerFallback(from, to string) *streamBuilder { + b.responses = append(b.responses, chat.MessageStreamResponse{ + Choices: []chat.MessageStreamChoice{{ + Index: 0, + Delta: chat.MessageDelta{ServerFallback: &chat.ServerFallback{From: from, To: to}}, + }}, + }) + return b +} + +func runServerFallbackStream(t *testing.T, stream *mockStream) (streamResult, []Event) { + t.Helper() + a := agent.New("root", "test", agent.WithModel(&mockProvider{id: "test/mock-model", stream: stream})) + sess := session.New(session.WithUserMessage("go")) + evCh := make(chan Event, 64) + res, err := handleStream(t.Context(), nil, stream, a, nil, sess, NewChannelSink(evCh), defaultStreamIdleTimeout) + require.NoError(t, err) + close(evCh) + var events []Event + for ev := range evCh { + events = append(events, ev) + } + return res, events +} + +// A server fallback supersedes the declined model's tool calls and thinking: +// they must never run or be replayed. Its text stays, as the API keeps it. +func TestHandleStream_ServerFallbackDropsSupersededWork(t *testing.T) { + t.Parallel() + + res, events := runServerFallbackStream(t, newStreamBuilder(). + AddReasoning("declined plan"). + AddContent("before "). + AddToolCallName("declined", "rm_rf"). + AddToolCallArguments("declined", `{"path":"/"}`). + AddServerFallback("claude-a", "claude-b"). + AddReasoning("new plan"). + AddContent("after"). + AddToolCallName("kept", "read"). + AddToolCallArguments("kept", `{}`). + AddToolCallStopWithUsage(1, 1). + Build()) + + require.Len(t, res.Calls, 1) + assert.Equal(t, "kept", res.Calls[0].ID) + assert.Equal(t, "before after", res.Content) + assert.Equal(t, "new plan", res.ReasoningContent) + + var warnings []string + for _, ev := range events { + if w, ok := ev.(*WarningEvent); ok { + warnings = append(warnings, w.Message) + } + } + require.Len(t, warnings, 1) + assert.Contains(t, warnings[0], "claude-a") + assert.Contains(t, warnings[0], "claude-b") +} + +func TestHandleStream_ServerFallbackWithoutSupersededWork(t *testing.T) { + t.Parallel() + + res, _ := runServerFallbackStream(t, newStreamBuilder(). + AddServerFallback("claude-a", "claude-b"). + AddContent("hi"). + AddStopWithUsage(1, 1). + Build()) + + assert.Empty(t, res.Calls) + assert.Equal(t, "hi", res.Content) +} diff --git a/pkg/runtime/streaming.go b/pkg/runtime/streaming.go index 1b69e27e50..95bee4bb9a 100644 --- a/pkg/runtime/streaming.go +++ b/pkg/runtime/streaming.go @@ -280,6 +280,23 @@ mainLoop: providerState = choice.Delta.ProviderState } + // A server fallback supersedes the declined model's tool calls and + // thinking: they must never run, and the next request drops them. + // Its text stays, as the API keeps it. + if fb := choice.Delta.ServerFallback; fb != nil { + responseStarted = true + if len(toolCalls) > 0 { + slog.WarnContext(ctx, "Dropping tool calls superseded by a server fallback", + "agent", a.Name(), "tool_calls", len(toolCalls), "from", fb.From, "to", fb.To) + } + toolCalls = nil + clear(toolCallIndex) + clear(emittedPartial) + fullReasoningContent.Reset() + thinkingSignature = "" + events.Emit(Warning(fmt.Sprintf("%s declined this turn; %s continues it.", fb.From, fb.To), a.Name())) + } + // A terminal chunk can also carry media; collect it before returning. if len(choice.Delta.Media) > 0 { responseStarted = true From 5ea034d233a71a10c332da00db10f1851d85885e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Wed, 21 Oct 2026 18:58:15 +0200 Subject: [PATCH 04/11] feat: retract tool calls a refusal or server fallback supersedes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The runtime already drops those calls, but the UIs kept showing them as pending. Emit tool_calls_retracted with their IDs; the TUI removes them, inside a reasoning block or standalone, and so does the lean TUI. Signed-off-by: Arnaud Héritier --- docs/features/api-server/index.md | 1 + pkg/leantui/events.go | 4 +++ pkg/leantui/retract_test.go | 24 ++++++++++++++ pkg/runtime/client.go | 1 + pkg/runtime/event.go | 18 ++++++++++ pkg/runtime/server_fallback_test.go | 28 ++++++++++++++-- pkg/runtime/streaming.go | 10 ++++++ .../components/messages/message_group_test.go | 33 +++++++++++++++++++ pkg/tui/components/messages/messages.go | 30 +++++++++++++++-- .../reasoningblock/reasoningblock.go | 27 +++++++++++++++ pkg/tui/page/chat/runtime_events.go | 3 ++ 11 files changed, 175 insertions(+), 4 deletions(-) create mode 100644 pkg/leantui/retract_test.go diff --git a/docs/features/api-server/index.md b/docs/features/api-server/index.md index 963772bc1f..a28c35414d 100644 --- a/docs/features/api-server/index.md +++ b/docs/features/api-server/index.md @@ -212,6 +212,7 @@ Event types include: - `agent_choice` — Streamed text content (partial responses) - `agent_choice_reasoning` — Streamed reasoning content - `partial_tool_call` — Streamed tool-call arguments, keyed by `tool_call.id`. The first one also carries `tool_definition` +- `tool_calls_retracted` — Tool calls announced by `partial_tool_call` that will never run, listed in `tool_call_ids`: the model refused, or a server fallback handed the turn to another model - `tool_call` — Agent requesting tool execution - `tool_call_output` — Streamed output of a running tool, keyed by `tool_call_id` - `tool_call_confirmation` — Tool call waiting for user approval diff --git a/pkg/leantui/events.go b/pkg/leantui/events.go index 6e11a08153..15288c6748 100644 --- a/pkg/leantui/events.go +++ b/pkg/leantui/events.go @@ -87,6 +87,10 @@ func (m *model) handleEvent(ctx context.Context, ev any) { identity := m.contentIdentity.Resolve(sessionID, e.Message.Message.MessageID) m.screen.Transcript.ReconcileAssistantContent(identity, chat.VisibleAssistantContent(e.Message.Message.Content)) m.contentIdentity.Finish(sessionID) + case *runtime.ToolCallsRetractedEvent: + for _, id := range e.ToolCallIDs { + m.screen.Transcript.RemoveTool(id) + } case *runtime.PartialToolCallEvent: m.screen.Transcript.FlushPending() toolDef := tools.Tool{Name: e.ToolCall.Function.Name} diff --git a/pkg/leantui/retract_test.go b/pkg/leantui/retract_test.go new file mode 100644 index 0000000000..230da9a3b6 --- /dev/null +++ b/pkg/leantui/retract_test.go @@ -0,0 +1,24 @@ +package leantui + +import ( + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/docker/docker-agent/pkg/leantui/ui" + "github.com/docker/docker-agent/pkg/runtime" + "github.com/docker/docker-agent/pkg/tools" +) + +func TestRetractedToolCallsAreRemoved(t *testing.T) { + t.Parallel() + m := &model{screen: ui.NewScreen("", "", "")} + for _, id := range []string{"declined", "kept"} { + m.handleEvent(t.Context(), runtime.PartialToolCall(tools.ToolCall{ID: id, Function: tools.FunctionCall{Name: "shell"}}, tools.Tool{}, "root")) + } + + m.handleEvent(t.Context(), runtime.ToolCallsRetracted([]string{"declined"}, "root")) + + assert.Nil(t, m.screen.Transcript.Tool("declined")) + assert.NotNil(t, m.screen.Transcript.Tool("kept")) +} diff --git a/pkg/runtime/client.go b/pkg/runtime/client.go index ce537d9e31..b09832df98 100644 --- a/pkg/runtime/client.go +++ b/pkg/runtime/client.go @@ -98,6 +98,7 @@ func NewClient(baseURL string, opts ...ClientOption) (*Client, error) { "session_summary": func() Event { return &SessionSummaryEvent{} }, "session_compaction": func() Event { return &SessionCompactionEvent{} }, "partial_tool_call": func() Event { return &PartialToolCallEvent{} }, + "tool_calls_retracted": func() Event { return &ToolCallsRetractedEvent{} }, "max_iterations_reached": func() Event { return &MaxIterationsReachedEvent{} }, "budget_usage": func() Event { return &BudgetUsageEvent{} }, "budget_exceeded": func() Event { return &BudgetExceededEvent{} }, diff --git a/pkg/runtime/event.go b/pkg/runtime/event.go index c319309f49..d497102822 100644 --- a/pkg/runtime/event.go +++ b/pkg/runtime/event.go @@ -100,6 +100,24 @@ func PartialToolCall(toolCall tools.ToolCall, toolDefinition tools.Tool, agentNa } } +// ToolCallsRetractedEvent is sent when tool calls announced by +// PartialToolCallEvent will never run: a refusal or a server fallback +// superseded them. +type ToolCallsRetractedEvent struct { + AgentContext + + Type string `json:"type"` + ToolCallIDs []string `json:"tool_call_ids"` +} + +func ToolCallsRetracted(toolCallIDs []string, agentName string) Event { + return &ToolCallsRetractedEvent{ + Type: "tool_calls_retracted", + ToolCallIDs: toolCallIDs, + AgentContext: newAgentContext(agentName), + } +} + // ToolCallEvent is sent when a tool call is received type ToolCallEvent struct { AgentContext diff --git a/pkg/runtime/server_fallback_test.go b/pkg/runtime/server_fallback_test.go index 42181b390e..7f2d36b790 100644 --- a/pkg/runtime/server_fallback_test.go +++ b/pkg/runtime/server_fallback_test.go @@ -60,11 +60,16 @@ func TestHandleStream_ServerFallbackDropsSupersededWork(t *testing.T) { assert.Equal(t, "new plan", res.ReasoningContent) var warnings []string + var retracted [][]string for _, ev := range events { - if w, ok := ev.(*WarningEvent); ok { - warnings = append(warnings, w.Message) + switch ev := ev.(type) { + case *WarningEvent: + warnings = append(warnings, ev.Message) + case *ToolCallsRetractedEvent: + retracted = append(retracted, ev.ToolCallIDs) } } + assert.Equal(t, [][]string{{"declined"}}, retracted, "the UI drops the pending call") require.Len(t, warnings, 1) assert.Contains(t, warnings[0], "claude-a") assert.Contains(t, warnings[0], "claude-b") @@ -82,3 +87,22 @@ func TestHandleStream_ServerFallbackWithoutSupersededWork(t *testing.T) { assert.Empty(t, res.Calls) assert.Equal(t, "hi", res.Content) } + +// A refusal voids the tool calls already shown to the user. +func TestHandleStream_RefusalRetractsToolCalls(t *testing.T) { + t.Parallel() + + _, events := runServerFallbackStream(t, newStreamBuilder(). + AddToolCallName("call_1", "rm_rf"). + AddToolCallArguments("call_1", `{"path":"/"}`). + AddRefusal(). + Build()) + + var retracted []string + for _, ev := range events { + if ev, ok := ev.(*ToolCallsRetractedEvent); ok { + retracted = append(retracted, ev.ToolCallIDs...) + } + } + assert.Equal(t, []string{"call_1"}, retracted) +} diff --git a/pkg/runtime/streaming.go b/pkg/runtime/streaming.go index 95bee4bb9a..0e707179a3 100644 --- a/pkg/runtime/streaming.go +++ b/pkg/runtime/streaming.go @@ -288,6 +288,7 @@ mainLoop: if len(toolCalls) > 0 { slog.WarnContext(ctx, "Dropping tool calls superseded by a server fallback", "agent", a.Name(), "tool_calls", len(toolCalls), "from", fb.From, "to", fb.To) + events.Emit(ToolCallsRetracted(toolCallIDs(toolCalls), a.Name())) } toolCalls = nil clear(toolCallIndex) @@ -386,6 +387,7 @@ mainLoop: if len(toolCalls) > 0 { slog.WarnContext(ctx, "Dropping tool calls from refused turn", "agent", a.Name(), "tool_calls", len(toolCalls)) + events.Emit(ToolCallsRetracted(toolCallIDs(toolCalls), a.Name())) toolCalls = nil openAIResponse = nil } @@ -513,3 +515,11 @@ mainLoop: ResponseStarted: responseStarted, }, nil } + +func toolCallIDs(calls []tools.ToolCall) []string { + ids := make([]string, len(calls)) + for i, tc := range calls { + ids[i] = tc.ID + } + return ids +} diff --git a/pkg/tui/components/messages/message_group_test.go b/pkg/tui/components/messages/message_group_test.go index 22804c8a07..df81af70a6 100644 --- a/pkg/tui/components/messages/message_group_test.go +++ b/pkg/tui/components/messages/message_group_test.go @@ -4,12 +4,14 @@ import ( "strconv" "testing" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/docker/docker-agent/pkg/runtime" "github.com/docker/docker-agent/pkg/session" "github.com/docker/docker-agent/pkg/tools" "github.com/docker/docker-agent/pkg/tui/animation" + "github.com/docker/docker-agent/pkg/tui/components/reasoningblock" "github.com/docker/docker-agent/pkg/tui/service" "github.com/docker/docker-agent/pkg/tui/types" ) @@ -132,3 +134,34 @@ func TestMessageGroupRebasedAfterLoadingRemoved(t *testing.T) { require.Len(t, m.messages, 3) require.Equal(t, "new reply", m.messages[2].Content) } + +// Retracted tool calls disappear wherever they render, and only them. +func TestRemoveToolCalls(t *testing.T) { + t.Parallel() + + m := NewScrollableView(animation.NewRuntime(), 80, 24, &service.SessionState{}).(*model) + m.AppendReasoning("root", "plan") + m.AddOrUpdateToolCall("root", tools.ToolCall{ID: "in-block", Function: tools.FunctionCall{Name: "shell"}}, tools.Tool{}, types.ToolStatusPending) + m.AddOrUpdateToolCall("root", tools.ToolCall{ID: "kept-in-block", Function: tools.FunctionCall{Name: "shell"}}, tools.Tool{}, types.ToolStatusPending) + m.BreakMessageGroup() + m.AddOrUpdateToolCall("root", tools.ToolCall{ID: "standalone", Function: tools.FunctionCall{Name: "shell"}}, tools.Tool{}, types.ToolStatusPending) + m.AddOrUpdateToolCall("root", tools.ToolCall{ID: "kept", Function: tools.FunctionCall{Name: "shell"}}, tools.Tool{}, types.ToolStatusPending) + + block, ok := m.views[0].(*reasoningblock.Model) + require.True(t, ok) + require.Equal(t, 2, block.ToolCount()) + + m.RemoveToolCalls([]string{"in-block", "standalone"}) + + assert.False(t, block.HasToolCall("in-block")) + assert.True(t, block.HasToolCall("kept-in-block")) + assert.Equal(t, 1, block.ToolCount()) + var standalone []string + for _, msg := range m.messages { + if msg.Type == types.MessageTypeToolCall { + standalone = append(standalone, msg.ToolCall.ID) + } + } + assert.Equal(t, []string{"kept"}, standalone) + assert.NotPanics(t, func() { _ = block.View() }) +} diff --git a/pkg/tui/components/messages/messages.go b/pkg/tui/components/messages/messages.go index 0a050686e7..4c93bbb511 100644 --- a/pkg/tui/components/messages/messages.go +++ b/pkg/tui/components/messages/messages.go @@ -86,6 +86,7 @@ type Model interface { AddCancelledMessage() tea.Cmd AddWelcomeMessage(content string) tea.Cmd AddOrUpdateToolCall(agentName string, toolCall tools.ToolCall, toolDef tools.Tool, status types.ToolStatus) tea.Cmd + RemoveToolCalls(toolCallIDs []string) AppendToolOutput(msg *runtime.ToolCallOutputEvent) tea.Cmd AddToolResult(msg *runtime.ToolCallResponseEvent, status types.ToolStatus) tea.Cmd AppendToLastMessage(agentName, content string) tea.Cmd @@ -2439,13 +2440,38 @@ func (m *model) removeSpinner() { } func (m *model) removePendingToolCallMessages() { + m.removeToolCallMessages(func(msg *types.Message) bool { + return msg.ToolStatus == types.ToolStatusPending || msg.ToolStatus == types.ToolStatusRunning + }) +} + +// RemoveToolCalls drops tool calls that will never run, whether they render +// on their own or inside a reasoning block. +func (m *model) RemoveToolCalls(toolCallIDs []string) { + for i, msg := range m.messages { + if msg.Type != types.MessageTypeAssistantReasoningBlock || i >= len(m.views) { + continue + } + if block, ok := m.views[i].(*reasoningblock.Model); ok { + for _, id := range toolCallIDs { + if block.RemoveToolCall(id) { + m.invalidateItem(i) + } + } + } + } + m.removeToolCallMessages(func(msg *types.Message) bool { + return slices.Contains(toolCallIDs, msg.ToolCall.ID) + }) +} + +func (m *model) removeToolCallMessages(match func(*types.Message) bool) { toolCallMessages := make([]*types.Message, 0, len(m.messages)) views := make([]layout.Model, 0, len(m.views)) removedBeforeGroup := 0 for i, msg := range m.messages { - if msg.Type == types.MessageTypeToolCall && - (msg.ToolStatus == types.ToolStatusPending || msg.ToolStatus == types.ToolStatusRunning) { + if msg.Type == types.MessageTypeToolCall && match(msg) { if i < m.groupStart { removedBeforeGroup++ } diff --git a/pkg/tui/components/reasoningblock/reasoningblock.go b/pkg/tui/components/reasoningblock/reasoningblock.go index 20555340d5..b0c1ce2990 100644 --- a/pkg/tui/components/reasoningblock/reasoningblock.go +++ b/pkg/tui/components/reasoningblock/reasoningblock.go @@ -3,6 +3,7 @@ package reasoningblock import ( "fmt" "math" + "slices" "strconv" "strings" "sync" @@ -338,6 +339,32 @@ func (m *Model) UpdateToolResult(toolCallID, content string, status types.ToolSt return nil } +// RemoveToolCall drops a tool call that will never run. It reports whether +// the block held it. +func (m *Model) RemoveToolCall(toolCallID string) bool { + k := slices.IndexFunc(m.toolEntries, func(e toolEntry) bool { return e.msg.ToolCall.ID == toolCallID }) + if k < 0 { + return false + } + animation.StopView(m.toolEntries[k].view) + m.toolEntries = slices.Delete(m.toolEntries, k, k+1) + items := m.contentItems[:0] + for _, item := range m.contentItems { + if item.kind == contentItemTool { + if item.toolIndex == k { + continue + } + if item.toolIndex > k { + item.toolIndex-- + } + } + items = append(items, item) + } + m.contentItems = items + m.cache = nil + return true +} + // HasToolCall returns true if the block contains the given tool call ID. func (m *Model) HasToolCall(toolCallID string) bool { for _, entry := range m.toolEntries { diff --git a/pkg/tui/page/chat/runtime_events.go b/pkg/tui/page/chat/runtime_events.go index cd6a4bbf69..bf7f72c206 100644 --- a/pkg/tui/page/chat/runtime_events.go +++ b/pkg/tui/page/chat/runtime_events.go @@ -118,6 +118,9 @@ func (p *chatPage) handleRuntimeEvent(msg tea.Msg) (bool, tea.Cmd) { return true, p.handleShellCommand(msg) // ===== Tool Events ===== + case *runtime.ToolCallsRetractedEvent: + p.messages.RemoveToolCalls(msg.ToolCallIDs) + return true, nil case *runtime.PartialToolCallEvent: return true, p.handlePartialToolCall(msg) From 7a4f14cef4581231f33fa6e661852e1ac497d0ab Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Wed, 21 Oct 2026 19:00:32 +0200 Subject: [PATCH 05/11] feat(anthropic): replay server fallback turns as the docs require MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Project a recorded turn before replay: keep the fallback markers, text, and everything after the final marker; drop thinking, redacted_thinking, connector_text and client tool_use before it, and server_tool_use without its result. Raw history stays intact, so a reload replays the same projection. Requests without the fallback beta, including the standard Messages API, drop the markers too: no thinking is left before them. Signed-off-by: Arnaud Héritier --- .../provider/anthropic/beta_converter.go | 5 +- .../anthropic/conversation_updates_test.go | 2 +- .../provider/anthropic/response_content.go | 82 +++++++++++-- .../anthropic/response_content_test.go | 8 +- .../anthropic/server_fallback_test.go | 108 ++++++++++++++++++ 5 files changed, 188 insertions(+), 17 deletions(-) diff --git a/pkg/model/provider/anthropic/beta_converter.go b/pkg/model/provider/anthropic/beta_converter.go index 16c9ea40e6..083857b56d 100644 --- a/pkg/model/provider/anthropic/beta_converter.go +++ b/pkg/model/provider/anthropic/beta_converter.go @@ -30,6 +30,7 @@ func (c *Client) convertBetaMessagesWithDeferred(ctx context.Context, messages [ var betaMessages []anthropic.BetaMessageParam deferredByCallID := deferredToolNamesByCallID(requestTools) loopStart := chat.ActiveToolLoopStart(messages) + keepMarkers := fallbacksEnabled(c.ModelConfig.ProviderOpts) for i := 0; i < len(messages); i++ { msg := &messages[i] @@ -70,13 +71,13 @@ func (c *Client) convertBetaMessagesWithDeferred(ctx context.Context, messages [ continue } if msg.Role == chat.MessageRoleAssistant { - if blocks, ok, err := betaReplayContent(msg); err != nil { + if blocks, ok, err := betaReplayContent(msg, keepMarkers); err != nil { return nil, err } else if ok { betaMessages = append(betaMessages, anthropic.BetaMessageParam{Role: anthropic.BetaMessageParamRoleAssistant, Content: blocks}) continue } - if blocks, ok := betaRecoveredParams(msg); ok { + if blocks, ok := betaRecoveredParams(msg, keepMarkers); ok { betaMessages = append(betaMessages, anthropic.BetaMessageParam{Role: anthropic.BetaMessageParamRoleAssistant, Content: blocks}) continue } diff --git a/pkg/model/provider/anthropic/conversation_updates_test.go b/pkg/model/provider/anthropic/conversation_updates_test.go index e89ffc1cf5..8b751e8fba 100644 --- a/pkg/model/provider/anthropic/conversation_updates_test.go +++ b/pkg/model/provider/anthropic/conversation_updates_test.go @@ -1027,7 +1027,7 @@ func TestWithRequestContext(t *testing.T) { assert.Equal(t, ctx, fromNil.RequestContext) msg := chat.Message{Role: chat.MessageRoleAssistant, Content: "x"} msg.AttachProviderState(fromNil) - params, ok, err := betaReplayContent(&msg) + params, ok, err := betaReplayContent(&msg, false) require.NoError(t, err) assert.False(t, ok, "a content-less state falls back to the flattened fields") assert.Nil(t, params) diff --git a/pkg/model/provider/anthropic/response_content.go b/pkg/model/provider/anthropic/response_content.go index af0585dcc6..94021d51fe 100644 --- a/pkg/model/provider/anthropic/response_content.go +++ b/pkg/model/provider/anthropic/response_content.go @@ -56,7 +56,7 @@ func replayContent(msg *chat.Message) ([]anthropic.ContentBlockParamUnion, bool, return nil, false, nil } var blocks []anthropic.ContentBlockUnion - if err := json.Unmarshal(state.Content, &blocks); err != nil { + if err := json.Unmarshal(projectServerFallback(state.Content, false), &blocks); err != nil { return nil, false, fmt.Errorf("anthropic: decoding raw response content: %w", err) } views := make([]rawBlock, len(blocks)) @@ -82,14 +82,15 @@ func replayContent(msg *chat.Message) ([]anthropic.ContentBlockParamUnion, bool, return params, true, nil } -// betaReplayContent is replayContent for the Beta Messages API. -func betaReplayContent(msg *chat.Message) ([]anthropic.BetaContentBlockParamUnion, bool, error) { +// betaReplayContent is replayContent for the Beta Messages API. keepMarkers +// keeps server fallback markers, for requests that send the fallback beta. +func betaReplayContent(msg *chat.Message, keepMarkers bool) ([]anthropic.BetaContentBlockParamUnion, bool, error) { state := msg.ReplayableProviderState(providerStateName) if state == nil { return nil, false, nil } var blocks []anthropic.BetaContentBlockUnion - if err := json.Unmarshal(state.Content, &blocks); err != nil { + if err := json.Unmarshal(projectServerFallback(state.Content, keepMarkers), &blocks); err != nil { return nil, false, fmt.Errorf("anthropic: decoding raw response content: %w", err) } views := make([]rawBlock, len(blocks)) @@ -181,12 +182,12 @@ func toolInput(raw string) any { // only those, so the signed and redacted thinking blocks stay valid. ok is // false when the thinking changed, the tool calls no longer line up, or the // state is from a format this build cannot read. -func recoveredContent(msg *chat.Message) (json.RawMessage, bool) { +func recoveredContent(msg *chat.Message, keepMarkers bool) (json.RawMessage, bool) { state := msg.ProviderState if !msg.HasOwnRawContent(providerStateName) || state.Version > chat.ProviderStateVersion { return nil, false } - decoder := json.NewDecoder(bytes.NewReader(state.Content)) + decoder := json.NewDecoder(bytes.NewReader(projectServerFallback(state.Content, keepMarkers))) decoder.UseNumber() var blocks []map[string]any if decoder.Decode(&blocks) != nil { @@ -244,7 +245,7 @@ func recoveredContent(msg *chat.Message) (json.RawMessage, bool) { // recoveredParams is recoveredContent decoded for the Messages API. func recoveredParams(msg *chat.Message) ([]anthropic.ContentBlockParamUnion, bool) { - raw, ok := recoveredContent(msg) + raw, ok := recoveredContent(msg, false) if !ok { return nil, false } @@ -266,8 +267,8 @@ func recoveredParams(msg *chat.Message) ([]anthropic.ContentBlockParamUnion, boo } // betaRecoveredParams is recoveredContent decoded for the Beta Messages API. -func betaRecoveredParams(msg *chat.Message) ([]anthropic.BetaContentBlockParamUnion, bool) { - raw, ok := recoveredContent(msg) +func betaRecoveredParams(msg *chat.Message, keepMarkers bool) ([]anthropic.BetaContentBlockParamUnion, bool) { + raw, ok := recoveredContent(msg, keepMarkers) if !ok { return nil, false } @@ -297,7 +298,7 @@ func recordedThinking(msg *chat.Message) bool { return true } var blocks []rawBlock - if json.Unmarshal(state.Content, &blocks) != nil { + if json.Unmarshal(projectServerFallback(state.Content, false), &blocks) != nil { return true } return slices.ContainsFunc(blocks, func(b rawBlock) bool { @@ -323,3 +324,64 @@ func (c *Client) flattenedThinking(i, loopStart int, msg *chat.Message) (bool, e slog.Warn("Anthropic: omitting thinking that changed after it was recorded", "index", i) return false, nil } + +// projectServerFallback returns a recorded response as the next request must +// send it after a server fallback: the fallback markers, text, and every block +// after the final marker stay; thinking, redacted_thinking, connector_text and +// client tool_use before it go, and so does a server_tool_use without its +// result. Kept blocks are byte-for-byte, so signatures stay valid. +// +// keepMarkers is false for requests without the fallback beta. The markers +// only anchor the thinking around them, and none is left before the final one. +func projectServerFallback(content json.RawMessage, keepMarkers bool) json.RawMessage { + var blocks []json.RawMessage + if json.Unmarshal(content, &blocks) != nil { + return content + } + type head struct { + Type string `json:"type"` + ID string `json:"id"` + ToolUseID string `json:"tool_use_id"` + } + heads := make([]head, len(blocks)) + final := -1 + results := map[string]bool{} + for i, b := range blocks { + _ = json.Unmarshal(b, &heads[i]) + if heads[i].Type == "fallback" { + final = i + } + if heads[i].ToolUseID != "" { + results[heads[i].ToolUseID] = true + } + } + if final < 0 { + return content + } + kept := make([]json.RawMessage, 0, len(blocks)) + for i, b := range blocks { + h := heads[i] + if h.Type == "fallback" { + if keepMarkers { + kept = append(kept, b) + } + continue + } + if i < final { + switch h.Type { + case "thinking", "redacted_thinking", "connector_text", "tool_use": + continue + case "server_tool_use": + if !results[h.ID] { + continue + } + } + } + kept = append(kept, b) + } + out, err := json.Marshal(kept) + if err != nil { + return content + } + return out +} diff --git a/pkg/model/provider/anthropic/response_content_test.go b/pkg/model/provider/anthropic/response_content_test.go index 35ffec21d7..0a1f66c78d 100644 --- a/pkg/model/provider/anthropic/response_content_test.go +++ b/pkg/model/provider/anthropic/response_content_test.go @@ -218,7 +218,7 @@ func TestReplayContent_ExactRoundTrip(t *testing.T) { require.NoError(t, err) assert.JSONEq(t, interleavedThinkingWire, string(got)) - betaBlocks, ok, err := betaReplayContent(&msg) + betaBlocks, ok, err := betaReplayContent(&msg, false) require.NoError(t, err) require.True(t, ok) got, err = json.Marshal(betaBlocks) @@ -281,7 +281,7 @@ func TestReplayContent_FallsBackWhenNotReplayable(t *testing.T) { assert.False(t, ok) assert.Nil(t, blocks) - betaBlocks, ok, err := betaReplayContent(&msg) + betaBlocks, ok, err := betaReplayContent(&msg, false) require.NoError(t, err) assert.False(t, ok) assert.Nil(t, betaBlocks) @@ -312,7 +312,7 @@ func TestReplayContent_FallsBackWhenRawDriftedBeforeSealing(t *testing.T) { _, ok, err := replayContent(&msg) require.NoError(t, err) assert.False(t, ok) - _, ok, err = betaReplayContent(&msg) + _, ok, err = betaReplayContent(&msg, false) require.NoError(t, err) assert.False(t, ok) }) @@ -360,7 +360,7 @@ func TestReplayContent_RejectsMalformedState(t *testing.T) { _, ok, err := replayContent(&msg) require.Error(t, err) assert.False(t, ok) - _, ok, err = betaReplayContent(&msg) + _, ok, err = betaReplayContent(&msg, false) require.Error(t, err) assert.False(t, ok) }) diff --git a/pkg/model/provider/anthropic/server_fallback_test.go b/pkg/model/provider/anthropic/server_fallback_test.go index d481a6353f..ac4eb34a8f 100644 --- a/pkg/model/provider/anthropic/server_fallback_test.go +++ b/pkg/model/provider/anthropic/server_fallback_test.go @@ -1,6 +1,7 @@ package anthropic import ( + "encoding/json" "errors" "io" "strings" @@ -13,6 +14,7 @@ import ( "github.com/docker/docker-agent/pkg/config/latest" "github.com/docker/docker-agent/pkg/model/provider/anthropic/anthropictest" "github.com/docker/docker-agent/pkg/model/provider/contracts" + "github.com/docker/docker-agent/pkg/tools" ) func fallbackClient(srv *scriptedServer) *Client { @@ -101,3 +103,109 @@ func TestServerFallbackStream(t *testing.T) { assert.Equal(t, to, model.Config.ModelConfig.Model) }) } + +// fallbackTurn is a turn the runtime recorded after a server fallback: the +// raw blocks are intact, while the flattened fields hold only the thinking +// and tool calls after the final marker, since the runtime dropped the rest. +func fallbackTurn(t *testing.T, raw, content, reasoning, signature string, calls ...string) chat.Message { + t.Helper() + msg := chat.Message{Role: chat.MessageRoleAssistant, Content: content, ReasoningContent: reasoning, ThinkingSignature: signature} + for _, id := range calls { + msg.ToolCalls = append(msg.ToolCalls, tools.ToolCall{ID: id, Type: "function", Function: tools.FunctionCall{Name: "read", Arguments: `{}`}}) + } + msg.AttachProviderState(&chat.ProviderState{Provider: providerStateName, Version: chat.ProviderStateVersion, Model: "claude-sonnet-4-6", Content: json.RawMessage(raw)}) + data, err := json.Marshal(msg) + require.NoError(t, err) + var reloaded chat.Message + require.NoError(t, json.Unmarshal(data, &reloaded)) + return reloaded +} + +const ( + fbMarker = `{"type":"fallback","from":{"model":"claude-opus-4-8"},"to":{"model":"claude-sonnet-4-6"},"trigger":{"type":"refusal","category":"cyber"}}` + fbMarker2 = `{"type":"fallback","from":{"model":"claude-sonnet-4-6"},"to":{"model":"claude-opus-4-7"},"trigger":{"type":"refusal","category":"cyber"}}` + fbDeclined = `{"type":"thinking","thinking":"declined","signature":"sig-d"}` + fbNew = `{"type":"thinking","thinking":"new","signature":"sig-n"}` +) + +func blockTypes(t *testing.T, v any) []string { + t.Helper() + data, err := json.Marshal(v) + require.NoError(t, err) + var blocks []map[string]any + require.NoError(t, json.Unmarshal(data, &blocks)) + types := make([]string, len(blocks)) + for i, b := range blocks { + types[i], _ = b["type"].(string) + if id, ok := b["id"].(string); ok { + types[i] += ":" + id + } + } + return types +} + +// The next request keeps the markers and the text, and drops what the +// declining models produced before the final marker, as the docs' continuation +// table requires. Without the fallback beta, the markers go too: no thinking +// is left before them, so nothing depends on their position. +func TestServerFallbackProjection(t *testing.T) { + t.Parallel() + + for name, tc := range map[string]struct { + msg chat.Message + withMarkers, plain []string + }{ + "mid output": { + msg: fallbackTurn(t, `[`+fbDeclined+`,{"type":"text","text":"before "},{"type":"tool_use","id":"t-declined","name":"read","input":{}},`+fbMarker+`,`+fbNew+`,{"type":"text","text":"after"},{"type":"tool_use","id":"t-kept","name":"read","input":{}}]`, + "before after", "new", "sig-n", "t-kept"), + withMarkers: []string{"text", "fallback", "thinking", "text", "tool_use:t-kept"}, + plain: []string{"text", "thinking", "text", "tool_use:t-kept"}, + }, + "two markers": { + msg: fallbackTurn(t, `[{"type":"text","text":"a "},`+fbMarker+`,{"type":"redacted_thinking","data":"r"},{"type":"text","text":"b "},`+fbMarker2+`,{"type":"text","text":"c"}]`, + "a b c", "", ""), + withMarkers: []string{"text", "fallback", "text", "fallback", "text"}, + plain: []string{"text", "text", "text"}, + }, + "server tools before the marker": { + msg: fallbackTurn(t, `[{"type":"server_tool_use","id":"s-paired","name":"web_search","input":{"query":"q"}},{"type":"web_search_tool_result","tool_use_id":"s-paired","content":[]},{"type":"server_tool_use","id":"s-open","name":"web_search","input":{"query":"q"}},`+fbMarker+`,{"type":"text","text":"answer"}]`, + "answer", "", ""), + withMarkers: []string{"server_tool_use:s-paired", "web_search_tool_result", "fallback", "text"}, + plain: []string{"server_tool_use:s-paired", "web_search_tool_result", "text"}, + }, + } { + t.Run(name, func(t *testing.T) { + t.Parallel() + beta, ok, err := betaReplayContent(&tc.msg, true) + require.NoError(t, err) + require.True(t, ok, "replayed raw, not rebuilt") + assert.Equal(t, tc.withMarkers, blockTypes(t, beta)) + + beta, ok, err = betaReplayContent(&tc.msg, false) + require.NoError(t, err) + require.True(t, ok) + assert.Equal(t, tc.plain, blockTypes(t, beta)) + }) + } + + t.Run("standard Messages API", func(t *testing.T) { + t.Parallel() + msg := fallbackTurn(t, `[`+fbDeclined+`,{"type":"text","text":"before "},`+fbMarker+`,{"type":"text","text":"after"}]`, "before after", "", "") + blocks, ok, err := replayContent(&msg) + require.NoError(t, err, "a marker must not reach the standard converter") + require.True(t, ok) + assert.Equal(t, []string{"text", "text"}, blockTypes(t, blocks)) + }) + + t.Run("signed thinking after the marker is replayed verbatim", func(t *testing.T) { + t.Parallel() + msg := fallbackTurn(t, `[`+fbDeclined+`,`+fbMarker+`,`+fbNew+`,{"type":"text","text":"x"}]`, "x", "new", "sig-n") + beta, ok, err := betaReplayContent(&msg, true) + require.NoError(t, err) + require.True(t, ok) + data, err := json.Marshal(beta) + require.NoError(t, err) + assert.Contains(t, string(data), `"signature":"sig-n"`) + assert.NotContains(t, string(data), "sig-d") + }) +} From 82970537ec78e1f3da1925456471ae125596c505 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Wed, 21 Oct 2026 19:02:44 +0200 Subject: [PATCH 06/11] test(runtime): a server fallback turn survives a reload MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A real Claude client streams a mid-output server fallback, then a sticky turn. The declined tool call never runs, both turns are stored under the fallback model, and later requests replay the projected turn, also after a SQLite reload. Signed-off-by: Arnaud Héritier --- .../claude5_compaction_lifecycle_test.go | 11 +- pkg/runtime/claude5_server_fallback_test.go | 101 ++++++++++++++++++ 2 files changed, 110 insertions(+), 2 deletions(-) create mode 100644 pkg/runtime/claude5_server_fallback_test.go diff --git a/pkg/runtime/claude5_compaction_lifecycle_test.go b/pkg/runtime/claude5_compaction_lifecycle_test.go index dc2b4e51e4..0b65e3adf6 100644 --- a/pkg/runtime/claude5_compaction_lifecycle_test.go +++ b/pkg/runtime/claude5_compaction_lifecycle_test.go @@ -27,6 +27,13 @@ import ( // claude5Runtime runs a real Claude Sonnet 5.5 client, built by the provider // registry, against server, with a 40k-token context window. func claude5Runtime(t *testing.T, server *anthropictest.Server, providerOpts map[string]any, store session.Store) *LocalRuntime { + t.Helper() + return anthropicRuntime(t, server, "claude-sonnet-5-5", providerOpts, store) +} + +// anthropicRuntime runs a real Anthropic client for model, built by the +// provider registry, against server, with a 40k-token context window. +func anthropicRuntime(t *testing.T, server *anthropictest.Server, modelName string, providerOpts map[string]any, store session.Store, agentOpts ...agent.Opt) *LocalRuntime { t.Helper() target, err := url.Parse(server.URL()) require.NoError(t, err) @@ -37,12 +44,12 @@ func claude5Runtime(t *testing.T, server *anthropictest.Server, providerOpts map return next.RoundTrip(req) }) }) - cfg := latest.ModelConfig{Provider: "anthropic", Model: "claude-sonnet-5-5", ProviderOpts: providerOpts} + cfg := latest.ModelConfig{Provider: "anthropic", Model: modelName, ProviderOpts: providerOpts} env := environment.NewMapEnvProvider(map[string]string{"ANTHROPIC_API_KEY": "test-key"}) model, err := providers.NewDefaultRegistry().New(t.Context(), &cfg, env, reroute) require.NoError(t, err) - root := agent.New("root", "You are terse.", agent.WithModel(model)) + root := agent.New("root", "You are terse.", append([]agent.Opt{agent.WithModel(model)}, agentOpts...)...) rt, err := NewLocalRuntime(t.Context(), team.New(team.WithAgents(root)), WithSessionStore(store), WithModelStore(mockModelStoreWithCostAndLimit{limit: 40_000, cost: modelsdev.Cost{Input: 1, Output: 2}})) diff --git a/pkg/runtime/claude5_server_fallback_test.go b/pkg/runtime/claude5_server_fallback_test.go new file mode 100644 index 0000000000..00f6b08829 --- /dev/null +++ b/pkg/runtime/claude5_server_fallback_test.go @@ -0,0 +1,101 @@ +package runtime + +import ( + "context" + "encoding/json" + "path/filepath" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/agent" + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/model/provider/anthropic/anthropictest" + "github.com/docker/docker-agent/pkg/session" + "github.com/docker/docker-agent/pkg/tools" +) + +// assistantBlocks returns the content block types of each assistant message +// in a captured request, with tool_use IDs and thinking signatures appended. +func assistantBlocks(t *testing.T, req anthropictest.Request) [][]string { + t.Helper() + messages, _ := req.Body["messages"].([]any) + var turns [][]string + for _, m := range messages { + msg, _ := m.(map[string]any) + if msg["role"] != "assistant" { + continue + } + content, _ := msg["content"].([]any) + var types []string + for _, c := range content { + block, _ := c.(map[string]any) + typ, _ := block["type"].(string) + switch typ { + case "tool_use": + typ += ":" + block["id"].(string) + case "thinking": + typ += ":" + block["signature"].(string) + } + types = append(types, typ) + } + turns = append(turns, types) + } + return turns +} + +// A server fallback mid-turn hands the turn to another model: the declined +// model's tool call never runs, the stored turn names the model that served +// it, and every later request, before and after a reload, replays the turn +// as the docs' continuation table requires. A later sticky-routed turn is +// attributed to the fallback model without any marker. +func TestServerFallbackLifecycle(t *testing.T) { + t.Parallel() + + const primary, fallback = "claude-opus-4-8", "claude-sonnet-4-6" + server := anthropictest.NewServer(t, + anthropictest.FallbackMidOutputReply("msg_1", primary, fallback, "before ", "after", "t-declined"), + anthropictest.StickyReply("msg_2", fallback, "sticky"), + anthropictest.TextReply("msg_3", "done", 10), + ) + var calls atomic.Int32 + record := agent.WithTools(tools.Tool{Name: "record", Parameters: map[string]any{"type": "object"}, Handler: func(context.Context, tools.ToolCall, tools.Runtime) (*tools.ToolCallResult, error) { + calls.Add(1) + return tools.ResultSuccess("recorded"), nil + }}) + opts := map[string]any{"fallbacks": []any{fallback}} + dbPath := filepath.Join(t.TempDir(), "sessions.db") + rt := anthropicRuntime(t, server, primary, opts, openSessionStore(t, dbPath), record) + + sess := session.New(session.WithUserMessage("u1"), session.WithToolsApproved(true)) + requireNoErrorEvents(t, drainRunStream(t, rt, sess)) + assert.Zero(t, calls.Load(), "the declined model's tool call must never run") + + sess.AddMessage(session.UserMessage("u2")) + requireNoErrorEvents(t, drainRunStream(t, rt, sess)) + + loaded, err := openSessionStore(t, dbPath).GetSession(t.Context(), sess.ID) + require.NoError(t, err) + var models []string + for _, item := range loaded.Messages { + if item.Message != nil && item.Message.Message.Role == chat.MessageRoleAssistant { + models = append(models, item.Message.Message.Model) + } + } + assert.Equal(t, []string{"anthropic/" + fallback, "anthropic/" + fallback}, models, "the marker and sticky routing both name the serving model") + + loaded.AddMessage(session.UserMessage("u3")) + reloaded := anthropicRuntime(t, server, primary, opts, openSessionStore(t, dbPath), record) + requireNoErrorEvents(t, drainRunStream(t, reloaded, loaded)) + + requests := server.Requests() + require.Len(t, requests, 3) + projected := []string{"text", "fallback", "text"} + assert.Equal(t, [][]string{projected}, assistantBlocks(t, requests[1]), "no declined thinking or tool_use") + assert.Equal(t, [][]string{projected, {"text"}}, assistantBlocks(t, requests[2]), "the same projection after a reload") + raw, _ := json.Marshal(requests[2].Body) + assert.NotContains(t, string(raw), "sig-declined") + assert.Zero(t, calls.Load()) +} From d912059a75afadc5468ed9f49a3a0d0f566da21c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Wed, 21 Oct 2026 19:03:18 +0200 Subject: [PATCH 07/11] test(anthropic): server fallback state stays per stream MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two concurrent streams on one client: only the one that fell back reports markers and a serving model. Signed-off-by: Arnaud Héritier --- .../anthropic/server_fallback_test.go | 69 +++++++++++++++++++ 1 file changed, 69 insertions(+) diff --git a/pkg/model/provider/anthropic/server_fallback_test.go b/pkg/model/provider/anthropic/server_fallback_test.go index ac4eb34a8f..3e40d07653 100644 --- a/pkg/model/provider/anthropic/server_fallback_test.go +++ b/pkg/model/provider/anthropic/server_fallback_test.go @@ -4,7 +4,10 @@ import ( "encoding/json" "errors" "io" + "net/http" + "net/http/httptest" "strings" + "sync" "testing" "github.com/stretchr/testify/assert" @@ -209,3 +212,69 @@ func TestServerFallbackProjection(t *testing.T) { assert.NotContains(t, string(data), "sig-d") }) } + +// Two concurrent streams on one client: the one that fell back reports its +// markers and serving model, the other reports neither. +func TestServerFallbackStatePerStream(t *testing.T) { + t.Parallel() + + const from, to = "claude-opus-4-8", "claude-sonnet-4-6" + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + raw, _ := io.ReadAll(r.Body) + w.Header().Set("content-type", "text/event-stream") + if strings.Contains(string(raw), "falls-back") { + _, _ = w.Write([]byte(anthropictest.FallbackMidOutputReply("m1", from, to, "before ", "after", "t1"))) + return + } + _, _ = w.Write([]byte(anthropictest.TextReply("m2", "plain", 10))) + })) + t.Cleanup(srv.Close) + client := fallbackClient(&scriptedServer{Server: srv}) + + type outcome struct { + markers []chat.ServerFallback + served string + err error + } + prompts := []string{"falls-back", "plain"} + results := make([]outcome, len(prompts)) + start := make(chan struct{}) + var wg sync.WaitGroup + for i, prompt := range prompts { + wg.Go(func() { + <-start + stream, err := client.CreateChatCompletionStream(t.Context(), []chat.Message{user(prompt)}, nil) + if err != nil { + results[i].err = err + return + } + defer stream.Close() + for { + resp, err := stream.Recv() + if err != nil { + if !errors.Is(err, io.EOF) { + results[i].err = err + } + break + } + for _, c := range resp.Choices { + if c.Delta.ServerFallback != nil { + results[i].markers = append(results[i].markers, *c.Delta.ServerFallback) + } + } + } + if model, ok := stream.(contracts.ServedStream).ServedBy(); ok { + results[i].served = model.Config.ModelConfig.Model + } + }) + } + close(start) + wg.Wait() + + require.NoError(t, results[0].err) + require.NoError(t, results[1].err) + assert.Equal(t, []chat.ServerFallback{{From: from, To: to}}, results[0].markers) + assert.Equal(t, to, results[0].served) + assert.Empty(t, results[1].markers) + assert.Empty(t, results[1].served) +} From ac91d59faf78e6c190716d8e94efa1a516290683 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Wed, 21 Oct 2026 19:13:02 +0200 Subject: [PATCH 08/11] fix: settle server fallback turns in the adapter and the wasm UI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - A tool call superseded by a marker no longer makes the turn a tool turn: the finish reason stays stop. - Sticky routing names the last fallback_message, the one that served. - The wasm front end forwards tool_calls_retracted. - Tests check the reset signature, the rendered reasoning block, and sticky routing among other iteration types. Signed-off-by: Arnaud Héritier --- cmd/wasm/README.md | 1 + cmd/wasm/events.go | 2 ++ cmd/wasm/events_test.go | 1 + cmd/wasm/main.go | 1 + pkg/model/provider/anthropic/beta_adapter.go | 12 ++++++--- .../anthropic/server_fallback_test.go | 27 +++++++++++++++++++ pkg/runtime/server_fallback_test.go | 9 +++++++ .../components/messages/message_group_test.go | 4 ++- 8 files changed, 53 insertions(+), 4 deletions(-) diff --git a/cmd/wasm/README.md b/cmd/wasm/README.md index e3940bd9dc..5c7d72ced2 100644 --- a/cmd/wasm/README.md +++ b/cmd/wasm/README.md @@ -265,6 +265,7 @@ cancels the previous). Safe to call before the call has started running. | --- | --- | | `delta` | `content?` or `reasoning?` — a streamed chunk. | | `tool_call_delta` | `id, name, arguments` — a streamed chunk of a tool call's arguments. | +| `tool_calls_retracted` | `ids` — tool calls announced by `tool_call_delta` that will never run (refusal or server fallback). | | `tool_call` | `id, name, args` — the call is about to run. | | `tool_confirmation` | `id, name, args` — the runtime waits for `confirm()`. | | `tool_output` | `id, name, output` — incremental output of a running call. | diff --git a/cmd/wasm/events.go b/cmd/wasm/events.go index 0404b276f4..2aa17ceeae 100644 --- a/cmd/wasm/events.go +++ b/cmd/wasm/events.go @@ -58,6 +58,8 @@ func (p *eventProjector) project(ev embeddedchat.Event) []map[string]any { "name": e.ToolCall.Function.Name, "arguments": e.ToolCall.Function.Arguments, }) + case *runtime.ToolCallsRetractedEvent: + return one(map[string]any{"type": "tool_calls_retracted", "ids": e.ToolCallIDs}) case *runtime.ToolCallEvent: return one(toolCallEvent("tool_call", e.ToolCall)) case *runtime.ToolCallConfirmationEvent: diff --git a/cmd/wasm/events_test.go b/cmd/wasm/events_test.go index 8a962b9333..760e2126f7 100644 --- a/cmd/wasm/events_test.go +++ b/cmd/wasm/events_test.go @@ -34,6 +34,7 @@ func TestProjectToolEvents(t *testing.T) { def := tools.Tool{Name: "echo"} assert.Equal(t, []map[string]any{{"type": "tool_call_delta", "id": "c1", "name": "echo", "arguments": `{"text":"x"}`}}, p.project(raw(runtime.PartialToolCall(call, def, "root")))) + assert.Equal(t, []map[string]any{{"type": "tool_calls_retracted", "ids": []string{"c1"}}}, p.project(raw(runtime.ToolCallsRetracted([]string{"c1"}, "root")))) assert.Equal(t, []map[string]any{{"type": "tool_call", "id": "c1", "name": "echo", "args": `{"text":"x"}`}}, p.project(raw(runtime.ToolCall(call, def, "root")))) assert.Equal(t, []map[string]any{{"type": "tool_confirmation", "id": "c1", "name": "echo", "args": `{"text":"x"}`}}, p.project(raw(runtime.ToolCallConfirmation(call, def, "root", nil)))) assert.Equal(t, []map[string]any{{"type": "tool_output", "id": "c1", "name": "echo", "output": "line"}}, p.project(raw(runtime.ToolCallOutput("c1", def, "line", "root")))) diff --git a/cmd/wasm/main.go b/cmd/wasm/main.go index 531d4adb3b..e87db0a225 100644 --- a/cmd/wasm/main.go +++ b/cmd/wasm/main.go @@ -54,6 +54,7 @@ // // {type: "delta", content?: string, reasoning?: string} // {type: "tool_call_delta", id, name, arguments} +// {type: "tool_calls_retracted", ids} // drop these deltas: they will never run // {type: "tool_call", id, name, args} // {type: "tool_confirmation", id, name, args} // answer with confirm() // {type: "tool_output", id, name, output} diff --git a/pkg/model/provider/anthropic/beta_adapter.go b/pkg/model/provider/anthropic/beta_adapter.go index cfa171789e..f46f53125f 100644 --- a/pkg/model/provider/anthropic/beta_adapter.go +++ b/pkg/model/provider/anthropic/beta_adapter.go @@ -4,6 +4,7 @@ import ( "encoding/json" "fmt" "log/slog" + "slices" "github.com/anthropics/anthropic-sdk-go" "github.com/anthropics/anthropic-sdk-go/packages/ssestream" @@ -85,6 +86,8 @@ func (a *betaStreamAdapter) Recv() (chat.MessageStreamResponse, error) { } response.Choices[0].Delta.ToolCalls = []tools.ToolCall{toolCall} case anthropic.BetaFallbackBlock: + // The declined model's tool calls are superseded. + a.toolCall = false a.served = block.To.Model response.Choices[0].Delta.ServerFallback = &chat.ServerFallback{From: block.From.Model, To: block.To.Model} case anthropic.BetaThinkingBlock: @@ -122,9 +125,12 @@ func (a *betaStreamAdapter) Recv() (chat.MessageStreamResponse, error) { a.stopReason = eventVariant.Delta.StopReason // Sticky routing sends no marker: only a fallback_message iteration // shows that a fallback model served the turn. - for _, it := range a.message.Usage.Iterations { - if it.Type == "fallback_message" && a.served == "" { - a.served = it.Model + if a.served == "" { + for _, it := range slices.Backward(a.message.Usage.Iterations) { + if it.Type == "fallback_message" { + a.served = it.Model + break + } } } if a.trackUsage { diff --git a/pkg/model/provider/anthropic/server_fallback_test.go b/pkg/model/provider/anthropic/server_fallback_test.go index 3e40d07653..e1f293e24a 100644 --- a/pkg/model/provider/anthropic/server_fallback_test.go +++ b/pkg/model/provider/anthropic/server_fallback_test.go @@ -68,6 +68,12 @@ func TestServerFallbackStream(t *testing.T) { markers: []chat.ServerFallback{{From: from, To: "claude-opus-4-7"}, {From: "claude-opus-4-7", To: to}}, served: to, }, "sticky routing": {reply: anthropictest.StickyReply("m", to, "hi"), served: to}, + // Other iteration types (compaction, advisor) don't name the server. + "sticky routing among other iterations": { + reply: strings.Replace(anthropictest.StickyReply("m", to, "hi"), `"iterations":[`, + `"iterations":[{"type":"compaction","model":"`+from+`","input_tokens":1,"output_tokens":1},`, 1), + served: to, + }, // message_start names "claude-test", not the dispatched model: a // response model alone is no fallback. "no fallback": {reply: anthropictest.TextReply("m", "hi", 10)}, @@ -92,6 +98,27 @@ func TestServerFallbackStream(t *testing.T) { }) } + t.Run("a superseded tool call doesn't make the turn a tool turn", func(t *testing.T) { + t.Parallel() + stream, err := fallbackClient(newScriptedServer(t, anthropictest.FallbackMidOutputReply("m", from, to, "before", "after", "t1"))).CreateChatCompletionStream(t.Context(), []chat.Message{user("q")}, nil) + require.NoError(t, err) + defer stream.Close() + var finish chat.FinishReason + for { + resp, err := stream.Recv() + if errors.Is(err, io.EOF) { + break + } + require.NoError(t, err) + for _, c := range resp.Choices { + if c.FinishReason != "" { + finish = c.FinishReason + } + } + } + assert.Equal(t, chat.FinishReasonStop, finish) + }) + t.Run("stream cut after the marker", func(t *testing.T) { t.Parallel() reply := anthropictest.FallbackMidOutputReply("m", from, to, "before", "after", "t1") diff --git a/pkg/runtime/server_fallback_test.go b/pkg/runtime/server_fallback_test.go index 7f2d36b790..043af729ec 100644 --- a/pkg/runtime/server_fallback_test.go +++ b/pkg/runtime/server_fallback_test.go @@ -43,6 +43,7 @@ func TestHandleStream_ServerFallbackDropsSupersededWork(t *testing.T) { res, events := runServerFallbackStream(t, newStreamBuilder(). AddReasoning("declined plan"). + AddThinkingSignature("sig-declined"). AddContent("before "). AddToolCallName("declined", "rm_rf"). AddToolCallArguments("declined", `{"path":"/"}`). @@ -58,6 +59,7 @@ func TestHandleStream_ServerFallbackDropsSupersededWork(t *testing.T) { assert.Equal(t, "kept", res.Calls[0].ID) assert.Equal(t, "before after", res.Content) assert.Equal(t, "new plan", res.ReasoningContent) + assert.Empty(t, res.ThinkingSignature, "the declined model's signature goes with its thinking") var warnings []string var retracted [][]string @@ -106,3 +108,10 @@ func TestHandleStream_RefusalRetractsToolCalls(t *testing.T) { } assert.Equal(t, []string{"call_1"}, retracted) } + +func (b *streamBuilder) AddThinkingSignature(signature string) *streamBuilder { + b.responses = append(b.responses, chat.MessageStreamResponse{ + Choices: []chat.MessageStreamChoice{{Index: 0, Delta: chat.MessageDelta{ThinkingSignature: signature}}}, + }) + return b +} diff --git a/pkg/tui/components/messages/message_group_test.go b/pkg/tui/components/messages/message_group_test.go index df81af70a6..ba2c05aad6 100644 --- a/pkg/tui/components/messages/message_group_test.go +++ b/pkg/tui/components/messages/message_group_test.go @@ -163,5 +163,7 @@ func TestRemoveToolCalls(t *testing.T) { } } assert.Equal(t, []string{"kept"}, standalone) - assert.NotPanics(t, func() { _ = block.View() }) + block.SetExpanded(true) + view := block.View() + assert.NotContains(t, view, "in-block", "the retracted call no longer renders") } From 1d9b8b631de77fb4ac0606f85b4c1bf764bc4601 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Wed, 21 Oct 2026 19:19:43 +0200 Subject: [PATCH 09/11] fix(anthropic): keep server fallback markers on every beta request MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The fallback-credit docs say a fallback block is accepted on any request without a beta header and must stay where it appeared. Keep the markers on the beta Messages API whether or not fallbacks is set; only the standard Messages API, which has no fallback block type, drops them. Signed-off-by: Arnaud Héritier --- pkg/model/provider/anthropic/beta_converter.go | 5 ++--- .../anthropic/conversation_updates_test.go | 2 +- .../provider/anthropic/response_content.go | 17 +++++++++-------- .../provider/anthropic/response_content_test.go | 8 ++++---- .../provider/anthropic/server_fallback_test.go | 12 ++++++------ 5 files changed, 22 insertions(+), 22 deletions(-) diff --git a/pkg/model/provider/anthropic/beta_converter.go b/pkg/model/provider/anthropic/beta_converter.go index 083857b56d..16c9ea40e6 100644 --- a/pkg/model/provider/anthropic/beta_converter.go +++ b/pkg/model/provider/anthropic/beta_converter.go @@ -30,7 +30,6 @@ func (c *Client) convertBetaMessagesWithDeferred(ctx context.Context, messages [ var betaMessages []anthropic.BetaMessageParam deferredByCallID := deferredToolNamesByCallID(requestTools) loopStart := chat.ActiveToolLoopStart(messages) - keepMarkers := fallbacksEnabled(c.ModelConfig.ProviderOpts) for i := 0; i < len(messages); i++ { msg := &messages[i] @@ -71,13 +70,13 @@ func (c *Client) convertBetaMessagesWithDeferred(ctx context.Context, messages [ continue } if msg.Role == chat.MessageRoleAssistant { - if blocks, ok, err := betaReplayContent(msg, keepMarkers); err != nil { + if blocks, ok, err := betaReplayContent(msg); err != nil { return nil, err } else if ok { betaMessages = append(betaMessages, anthropic.BetaMessageParam{Role: anthropic.BetaMessageParamRoleAssistant, Content: blocks}) continue } - if blocks, ok := betaRecoveredParams(msg, keepMarkers); ok { + if blocks, ok := betaRecoveredParams(msg); ok { betaMessages = append(betaMessages, anthropic.BetaMessageParam{Role: anthropic.BetaMessageParamRoleAssistant, Content: blocks}) continue } diff --git a/pkg/model/provider/anthropic/conversation_updates_test.go b/pkg/model/provider/anthropic/conversation_updates_test.go index 8b751e8fba..e89ffc1cf5 100644 --- a/pkg/model/provider/anthropic/conversation_updates_test.go +++ b/pkg/model/provider/anthropic/conversation_updates_test.go @@ -1027,7 +1027,7 @@ func TestWithRequestContext(t *testing.T) { assert.Equal(t, ctx, fromNil.RequestContext) msg := chat.Message{Role: chat.MessageRoleAssistant, Content: "x"} msg.AttachProviderState(fromNil) - params, ok, err := betaReplayContent(&msg, false) + params, ok, err := betaReplayContent(&msg) require.NoError(t, err) assert.False(t, ok, "a content-less state falls back to the flattened fields") assert.Nil(t, params) diff --git a/pkg/model/provider/anthropic/response_content.go b/pkg/model/provider/anthropic/response_content.go index 94021d51fe..87b24636a7 100644 --- a/pkg/model/provider/anthropic/response_content.go +++ b/pkg/model/provider/anthropic/response_content.go @@ -82,15 +82,15 @@ func replayContent(msg *chat.Message) ([]anthropic.ContentBlockParamUnion, bool, return params, true, nil } -// betaReplayContent is replayContent for the Beta Messages API. keepMarkers -// keeps server fallback markers, for requests that send the fallback beta. -func betaReplayContent(msg *chat.Message, keepMarkers bool) ([]anthropic.BetaContentBlockParamUnion, bool, error) { +// betaReplayContent is replayContent for the Beta Messages API, which keeps +// server fallback markers. +func betaReplayContent(msg *chat.Message) ([]anthropic.BetaContentBlockParamUnion, bool, error) { state := msg.ReplayableProviderState(providerStateName) if state == nil { return nil, false, nil } var blocks []anthropic.BetaContentBlockUnion - if err := json.Unmarshal(projectServerFallback(state.Content, keepMarkers), &blocks); err != nil { + if err := json.Unmarshal(projectServerFallback(state.Content, true), &blocks); err != nil { return nil, false, fmt.Errorf("anthropic: decoding raw response content: %w", err) } views := make([]rawBlock, len(blocks)) @@ -267,8 +267,8 @@ func recoveredParams(msg *chat.Message) ([]anthropic.ContentBlockParamUnion, boo } // betaRecoveredParams is recoveredContent decoded for the Beta Messages API. -func betaRecoveredParams(msg *chat.Message, keepMarkers bool) ([]anthropic.BetaContentBlockParamUnion, bool) { - raw, ok := recoveredContent(msg, keepMarkers) +func betaRecoveredParams(msg *chat.Message) ([]anthropic.BetaContentBlockParamUnion, bool) { + raw, ok := recoveredContent(msg, true) if !ok { return nil, false } @@ -331,8 +331,9 @@ func (c *Client) flattenedThinking(i, loopStart int, msg *chat.Message) (bool, e // client tool_use before it go, and so does a server_tool_use without its // result. Kept blocks are byte-for-byte, so signatures stay valid. // -// keepMarkers is false for requests without the fallback beta. The markers -// only anchor the thinking around them, and none is left before the final one. +// keepMarkers is false on the standard Messages API, whose types have no +// fallback block. Dropping the markers is safe there: they only anchor the +// thinking around them, and none is left before the final one. func projectServerFallback(content json.RawMessage, keepMarkers bool) json.RawMessage { var blocks []json.RawMessage if json.Unmarshal(content, &blocks) != nil { diff --git a/pkg/model/provider/anthropic/response_content_test.go b/pkg/model/provider/anthropic/response_content_test.go index 0a1f66c78d..35ffec21d7 100644 --- a/pkg/model/provider/anthropic/response_content_test.go +++ b/pkg/model/provider/anthropic/response_content_test.go @@ -218,7 +218,7 @@ func TestReplayContent_ExactRoundTrip(t *testing.T) { require.NoError(t, err) assert.JSONEq(t, interleavedThinkingWire, string(got)) - betaBlocks, ok, err := betaReplayContent(&msg, false) + betaBlocks, ok, err := betaReplayContent(&msg) require.NoError(t, err) require.True(t, ok) got, err = json.Marshal(betaBlocks) @@ -281,7 +281,7 @@ func TestReplayContent_FallsBackWhenNotReplayable(t *testing.T) { assert.False(t, ok) assert.Nil(t, blocks) - betaBlocks, ok, err := betaReplayContent(&msg, false) + betaBlocks, ok, err := betaReplayContent(&msg) require.NoError(t, err) assert.False(t, ok) assert.Nil(t, betaBlocks) @@ -312,7 +312,7 @@ func TestReplayContent_FallsBackWhenRawDriftedBeforeSealing(t *testing.T) { _, ok, err := replayContent(&msg) require.NoError(t, err) assert.False(t, ok) - _, ok, err = betaReplayContent(&msg, false) + _, ok, err = betaReplayContent(&msg) require.NoError(t, err) assert.False(t, ok) }) @@ -360,7 +360,7 @@ func TestReplayContent_RejectsMalformedState(t *testing.T) { _, ok, err := replayContent(&msg) require.Error(t, err) assert.False(t, ok) - _, ok, err = betaReplayContent(&msg, false) + _, ok, err = betaReplayContent(&msg) require.Error(t, err) assert.False(t, ok) }) diff --git a/pkg/model/provider/anthropic/server_fallback_test.go b/pkg/model/provider/anthropic/server_fallback_test.go index e1f293e24a..e91bd9e133 100644 --- a/pkg/model/provider/anthropic/server_fallback_test.go +++ b/pkg/model/provider/anthropic/server_fallback_test.go @@ -176,8 +176,8 @@ func blockTypes(t *testing.T, v any) []string { // The next request keeps the markers and the text, and drops what the // declining models produced before the final marker, as the docs' continuation -// table requires. Without the fallback beta, the markers go too: no thinking -// is left before them, so nothing depends on their position. +// table requires. The standard Messages API has no fallback block, so the +// markers go too: no thinking is left before them. func TestServerFallbackProjection(t *testing.T) { t.Parallel() @@ -206,15 +206,15 @@ func TestServerFallbackProjection(t *testing.T) { } { t.Run(name, func(t *testing.T) { t.Parallel() - beta, ok, err := betaReplayContent(&tc.msg, true) + beta, ok, err := betaReplayContent(&tc.msg) require.NoError(t, err) require.True(t, ok, "replayed raw, not rebuilt") assert.Equal(t, tc.withMarkers, blockTypes(t, beta)) - beta, ok, err = betaReplayContent(&tc.msg, false) + standard, ok, err := replayContent(&tc.msg) require.NoError(t, err) require.True(t, ok) - assert.Equal(t, tc.plain, blockTypes(t, beta)) + assert.Equal(t, tc.plain, blockTypes(t, standard)) }) } @@ -230,7 +230,7 @@ func TestServerFallbackProjection(t *testing.T) { t.Run("signed thinking after the marker is replayed verbatim", func(t *testing.T) { t.Parallel() msg := fallbackTurn(t, `[`+fbDeclined+`,`+fbMarker+`,`+fbNew+`,{"type":"text","text":"x"}]`, "x", "new", "sig-n") - beta, ok, err := betaReplayContent(&msg, true) + beta, ok, err := betaReplayContent(&msg) require.NoError(t, err) require.True(t, ok) data, err := json.Marshal(beta) From e491ec7ea6ea927091f792d97687ab6dc914f168 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Thu, 22 Oct 2026 11:03:20 +0200 Subject: [PATCH 10/11] fix: close three gaps in server fallback handling MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Don't mine the declined model's kept text for XML tool calls after a handoff, and let the fallback model's text reach the UI again. - Convert the retracted IDs for JavaScript in the wasm front end. - Store the served model atomically: the runtime can read it while the reader goroutine still streams, after a cancellation. Signed-off-by: Arnaud Héritier --- cmd/wasm/events.go | 2 +- cmd/wasm/events_test.go | 5 ++- pkg/model/provider/anthropic/beta_adapter.go | 18 ++++++----- .../anthropic/server_fallback_test.go | 31 +++++++++++++++++++ pkg/runtime/server_fallback_test.go | 23 ++++++++++++++ pkg/runtime/streaming.go | 7 ++++- 6 files changed, 76 insertions(+), 10 deletions(-) diff --git a/cmd/wasm/events.go b/cmd/wasm/events.go index 2aa17ceeae..2e7262e3a8 100644 --- a/cmd/wasm/events.go +++ b/cmd/wasm/events.go @@ -59,7 +59,7 @@ func (p *eventProjector) project(ev embeddedchat.Event) []map[string]any { "arguments": e.ToolCall.Function.Arguments, }) case *runtime.ToolCallsRetractedEvent: - return one(map[string]any{"type": "tool_calls_retracted", "ids": e.ToolCallIDs}) + return one(map[string]any{"type": "tool_calls_retracted", "ids": stringsToAny(e.ToolCallIDs)}) case *runtime.ToolCallEvent: return one(toolCallEvent("tool_call", e.ToolCall)) case *runtime.ToolCallConfirmationEvent: diff --git a/cmd/wasm/events_test.go b/cmd/wasm/events_test.go index 760e2126f7..0529189373 100644 --- a/cmd/wasm/events_test.go +++ b/cmd/wasm/events_test.go @@ -4,6 +4,7 @@ package main import ( "errors" + "syscall/js" "testing" "github.com/stretchr/testify/assert" @@ -34,7 +35,9 @@ func TestProjectToolEvents(t *testing.T) { def := tools.Tool{Name: "echo"} assert.Equal(t, []map[string]any{{"type": "tool_call_delta", "id": "c1", "name": "echo", "arguments": `{"text":"x"}`}}, p.project(raw(runtime.PartialToolCall(call, def, "root")))) - assert.Equal(t, []map[string]any{{"type": "tool_calls_retracted", "ids": []string{"c1"}}}, p.project(raw(runtime.ToolCallsRetracted([]string{"c1"}, "root")))) + retracted := p.project(raw(runtime.ToolCallsRetracted([]string{"c1"}, "root"))) + assert.Equal(t, []map[string]any{{"type": "tool_calls_retracted", "ids": []any{"c1"}}}, retracted) + assert.NotPanics(t, func() { js.ValueOf(retracted[0]) }, "emit converts the event for JavaScript") assert.Equal(t, []map[string]any{{"type": "tool_call", "id": "c1", "name": "echo", "args": `{"text":"x"}`}}, p.project(raw(runtime.ToolCall(call, def, "root")))) assert.Equal(t, []map[string]any{{"type": "tool_confirmation", "id": "c1", "name": "echo", "args": `{"text":"x"}`}}, p.project(raw(runtime.ToolCallConfirmation(call, def, "root", nil)))) assert.Equal(t, []map[string]any{{"type": "tool_output", "id": "c1", "name": "echo", "output": "line"}}, p.project(raw(runtime.ToolCallOutput("c1", def, "line", "root")))) diff --git a/pkg/model/provider/anthropic/beta_adapter.go b/pkg/model/provider/anthropic/beta_adapter.go index f46f53125f..bd20ca0666 100644 --- a/pkg/model/provider/anthropic/beta_adapter.go +++ b/pkg/model/provider/anthropic/beta_adapter.go @@ -5,6 +5,7 @@ import ( "fmt" "log/slog" "slices" + "sync/atomic" "github.com/anthropics/anthropic-sdk-go" "github.com/anthropics/anthropic-sdk-go/packages/ssestream" @@ -31,9 +32,11 @@ type betaStreamAdapter struct { rawLost bool requestContext json.RawMessage // config is the dispatching client's; served is the model a server - // fallback handed the turn to, empty when the dispatched model served it. + // fallback handed the turn to, unset when the dispatched model served it. + // The runtime may read it while the reader goroutine still writes it + // (cancellation, idle timeout). config base.Config - served string + served atomic.Pointer[string] } // newBetaStreamAdapter creates a new Beta stream adapter @@ -88,7 +91,7 @@ func (a *betaStreamAdapter) Recv() (chat.MessageStreamResponse, error) { case anthropic.BetaFallbackBlock: // The declined model's tool calls are superseded. a.toolCall = false - a.served = block.To.Model + a.served.Store(&block.To.Model) response.Choices[0].Delta.ServerFallback = &chat.ServerFallback{From: block.From.Model, To: block.To.Model} case anthropic.BetaThinkingBlock: if block.Thinking != "" { @@ -125,10 +128,10 @@ func (a *betaStreamAdapter) Recv() (chat.MessageStreamResponse, error) { a.stopReason = eventVariant.Delta.StopReason // Sticky routing sends no marker: only a fallback_message iteration // shows that a fallback model served the turn. - if a.served == "" { + if a.served.Load() == nil { for _, it := range slices.Backward(a.message.Usage.Iterations) { if it.Type == "fallback_message" { - a.served = it.Model + a.served.Store(&it.Model) break } } @@ -171,11 +174,12 @@ func (a *betaStreamAdapter) accumulate(event anthropic.BetaRawMessageStreamEvent // ServedBy reports the model a server fallback handed the turn to. func (a *betaStreamAdapter) ServedBy() (contracts.ServingModel, bool) { - if a.served == "" { + served := a.served.Load() + if served == nil { return contracts.ServingModel{}, false } cfg := a.config - cfg.ModelConfig.Model = a.served + cfg.ModelConfig.Model = *served // The primary's display name and cost overrides don't describe the // fallback model, which is priced from the catalogue. cfg.ModelConfig.DisplayModel = "" diff --git a/pkg/model/provider/anthropic/server_fallback_test.go b/pkg/model/provider/anthropic/server_fallback_test.go index e91bd9e133..aab311534d 100644 --- a/pkg/model/provider/anthropic/server_fallback_test.go +++ b/pkg/model/provider/anthropic/server_fallback_test.go @@ -305,3 +305,34 @@ func TestServerFallbackStatePerStream(t *testing.T) { assert.Empty(t, results[1].markers) assert.Empty(t, results[1].served) } + +// The runtime can read ServedBy while the reader goroutine still streams, after +// a cancellation or an idle timeout. +func TestServedByWhileStreaming(t *testing.T) { + t.Parallel() + + stream, err := fallbackClient(newScriptedServer(t, anthropictest.FallbackMidOutputReply("m", "claude-opus-4-8", "claude-sonnet-4-6", "before", "after", "t1"))).CreateChatCompletionStream(t.Context(), []chat.Message{user("q")}, nil) + require.NoError(t, err) + defer stream.Close() + served := stream.(contracts.ServedStream) + + done := make(chan struct{}) + go func() { + defer close(done) + for { + if _, err := stream.Recv(); err != nil { + return + } + } + }() + for { + served.ServedBy() + select { + case <-done: + _, ok := served.ServedBy() + assert.True(t, ok) + return + default: + } + } +} diff --git a/pkg/runtime/server_fallback_test.go b/pkg/runtime/server_fallback_test.go index 043af729ec..decc13307b 100644 --- a/pkg/runtime/server_fallback_test.go +++ b/pkg/runtime/server_fallback_test.go @@ -1,6 +1,7 @@ package runtime import ( + "strings" "testing" "github.com/stretchr/testify/assert" @@ -115,3 +116,25 @@ func (b *streamBuilder) AddThinkingSignature(signature string) *streamBuilder { }) return b } + +// Text kept from the declined model must not resurrect an XML tool call, and +// the fallback model's text still reaches the UI. +func TestHandleStream_ServerFallbackIgnoresDeclinedXMLToolCalls(t *testing.T) { + t.Parallel() + + res, events := runServerFallbackStream(t, newStreamBuilder(). + AddContent(`{"name":"rm_rf","arguments":{}}`). + AddServerFallback("claude-a", "claude-b"). + AddContent("answer"). + AddStopWithUsage(1, 1). + Build()) + + assert.Empty(t, res.Calls, "the declined model's XML call must never run") + var shown strings.Builder + for _, ev := range events { + if c, ok := ev.(*AgentChoiceEvent); ok { + shown.WriteString(c.Content) + } + } + assert.Contains(t, shown.String(), "answer") +} diff --git a/pkg/runtime/streaming.go b/pkg/runtime/streaming.go index 0e707179a3..44e9c779a6 100644 --- a/pkg/runtime/streaming.go +++ b/pkg/runtime/streaming.go @@ -155,6 +155,9 @@ func handleStream(ctx context.Context, cancelStream context.CancelCauseFunc, str // xmlToolCallGate suppresses AgentChoice events once a tag is // seen, preventing raw XML from rendering in the TUI. xmlToolCallGate := false + // serverHandoff records a server fallback: the text kept from the declined + // model must not be mined for XML tool calls it was never allowed to make. + serverHandoff := false for _, t := range agentTools { toolDefMap[t.Name] = t } @@ -199,7 +202,7 @@ func handleStream(ctx context.Context, cancelStream context.CancelCauseFunc, str // no structured tool calls were received. Called from both the early-return // and EOF paths. applyXMLFallback := func() { - if len(toolCalls) > 0 { + if len(toolCalls) > 0 || serverHandoff { return } extracted, textBefore, found := extractXMLToolCalls(fullContent.String()) @@ -295,6 +298,8 @@ mainLoop: clear(emittedPartial) fullReasoningContent.Reset() thinkingSignature = "" + serverHandoff = true + xmlToolCallGate = false events.Emit(Warning(fmt.Sprintf("%s declined this turn; %s continues it.", fb.From, fb.To), a.Name())) } From 36219fcb2af6f459a0581bae14803fcfd5ea5316 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Thu, 22 Oct 2026 13:51:16 +0200 Subject: [PATCH 11/11] fix(anthropic): size the recovered content without arithmetic MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CodeQL flags len(blocks)+1 as a possible allocation-size overflow now that the blocks come through json.Marshal. append grows the slice for the one extra text block when it is needed. Signed-off-by: Arnaud Héritier --- pkg/model/provider/anthropic/response_content.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/model/provider/anthropic/response_content.go b/pkg/model/provider/anthropic/response_content.go index 87b24636a7..539e738173 100644 --- a/pkg/model/provider/anthropic/response_content.go +++ b/pkg/model/provider/anthropic/response_content.go @@ -224,7 +224,7 @@ func recoveredContent(msg *chat.Message, keepMarkers bool) (json.RawMessage, boo drop[i] = true } } - patched := make([]map[string]any, 0, len(blocks)+1) + patched := make([]map[string]any, 0, len(blocks)) inserted := len(text) > 0 || msg.Content == "" for i, b := range blocks { if drop[i] {