From 8fafa7013066a12d135d2cf8c6606009a67b8ebf Mon Sep 17 00:00:00 2001 From: MK Date: Fri, 31 Jul 2026 19:17:53 -0400 Subject: [PATCH 1/3] fix(mcp): route call-time ErrNoToken through the auth gate (closes #376) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The delegated-consent auth gate (#330) was consulted only for resolveClient (connection-establish) errors. But per-request-auth transports (transport_http attaches the bearer on every Send, including tools/call) 403 at CALL time, not at initialize — so a type=user server whose user has no grant surfaced ErrNoToken from client.CallTool, which the gate never inspected. The call failed hard with reason=no_token, no mcp_auth_required event fired, and the platform consent flow had nothing to trigger on (field: Atlassian agent, forge#376). Execute now wraps the whole resolve→call sequence in the gate: an ErrNoToken from either half parks via authGate.Await and, on a granted resume, retries resolve→call. Same bounded, one-shot semantics as before; the diagnostic phase prefix is irrelevant now since both paths are covered. Tests: call-time no-token parks + retries to success; gate give-up still fails as no_token; a non-auth CallTool error never parks; nil gate still surfaces ErrNoToken. --- forge-core/tools/adapters/mcp_tool.go | 40 ++++++---- forge-core/tools/adapters/mcp_tool_test.go | 88 ++++++++++++++++++++++ 2 files changed, 112 insertions(+), 16 deletions(-) diff --git a/forge-core/tools/adapters/mcp_tool.go b/forge-core/tools/adapters/mcp_tool.go index daba571..e459f60 100644 --- a/forge-core/tools/adapters/mcp_tool.go +++ b/forge-core/tools/adapters/mcp_tool.go @@ -198,29 +198,37 @@ func (m *MCPTool) Execute(ctx context.Context, args json.RawMessage) (string, er correlationID := runtime.CorrelationIDFromContext(ctx) m.emitCall(correlationID, len(args)) - // Resolve the connection for THIS call's requesting user (#317). A - // static resolver returns the shared client (unchanged); a per-subject - // pool returns that user's own connection, establishing it lazily. - client, err := m.resolveClient(ctx) + // resolveAndCall runs the whole resolve→call sequence for THIS request. + // ErrNoToken can surface from EITHER half: the per-subject connection + // establish (transports that authenticate at initialize) OR CallTool + // itself (transports that attach the token per-request and only 403 on + // the tools/call frame — Atlassian and most OAuth MCP servers). The auth + // gate (#330) must catch both, so it wraps the whole sequence rather than + // just resolveClient (the call-time case previously slipped past the gate + // and failed hard with reason=no_token — forge#376). + resolveAndCall := func() (*mcp.CallToolResult, error) { + client, err := m.resolveClient(ctx) + if err != nil { + return nil, err + } + return client.CallTool(ctx, m.descriptor.Name, args) + } + + res, err := resolveAndCall() if err != nil && m.authGate != nil && errors.Is(err, mcp.ErrNoToken) { // No grant yet for this user (#330). Rather than fail the call, park - // the executor and let the user consent; on a granted resume, - // re-resolve — the delegated path now finds the grant and the - // per-user connection establishes. A gate error (timeout / cancel / - // no requesting user) means give up; it flows to the emit+return - // below and classifies like the underlying ErrNoToken. + // the executor and let the user consent; on a granted resume, retry + // the full resolve→call — the delegated path now finds the grant, the + // per-user connection establishes, and the tool call goes through. A + // gate error (timeout / cancel / no requesting user) means give up; + // it flows to the emit+return below and classifies like the + // underlying ErrNoToken. if gateErr := m.authGate.Await(ctx, m.server); gateErr != nil { err = gateErr } else { - client, err = m.resolveClient(ctx) + res, err = resolveAndCall() } } - if err != nil { - durMs := time.Since(start).Milliseconds() - m.emitResult(correlationID, durMs, 0, false, classifyToolErr(err)) - return "", fmt.Errorf("mcp %s/%s: %w", m.server, m.descriptor.Name, err) - } - res, err := client.CallTool(ctx, m.descriptor.Name, args) durMs := time.Since(start).Milliseconds() if err != nil { diff --git a/forge-core/tools/adapters/mcp_tool_test.go b/forge-core/tools/adapters/mcp_tool_test.go index d875ca9..744a448 100644 --- a/forge-core/tools/adapters/mcp_tool_test.go +++ b/forge-core/tools/adapters/mcp_tool_test.go @@ -343,3 +343,91 @@ func (b *safeBuf) String() string { defer b.mu.Store(0) return b.buf.String() } + +// gateStub records Await calls and returns a scripted outcome. nil err ⇒ +// "granted" (caller retries); non-nil ⇒ give up. +type gateStub struct { + calls atomic.Int32 + err error + onAwait func() // optional side effect on grant (e.g. flip the client to succeed) +} + +func (g *gateStub) Await(context.Context, string) error { + g.calls.Add(1) + if g.onAwait != nil { + g.onAwait() + } + return g.err +} + +// TestMCPTool_Execute_CallTimeNoToken_Parks is the forge#376 regression: an +// ErrNoToken raised by CallTool (per-request auth transports — the token is +// attached per frame, so the 403 lands on tools/call, not at establish) must +// route through the auth gate and retry, exactly like an establish-time +// ErrNoToken. Before the fix this path never consulted the gate and failed +// hard with reason=no_token. +func TestMCPTool_Execute_CallTimeNoToken_Parks(t *testing.T) { + t.Parallel() + // Client 403s on the first CallTool, succeeds after consent (the gate's + // onAwait flips it), modelling "grant now exists → retry resolves it". + c := &mockClient{err: mcp.ErrNoToken} + gate := &gateStub{onAwait: func() { + c.err = nil + c.res = &mcp.CallToolResult{Content: []mcp.ToolContent{{Type: "text", Text: "ok-after-consent"}}} + }} + a := newAdapter(t, c, func(m *MCPTool) { m.authGate = gate }) + + got, err := a.Execute(context.Background(), json.RawMessage(`{}`)) + if err != nil { + t.Fatalf("expected park+retry to succeed, got err=%v", err) + } + if gate.calls.Load() != 1 { + t.Fatalf("auth gate consulted %d times, want 1 (a call-time no-token must park)", gate.calls.Load()) + } + if got != "ok-after-consent" { + t.Fatalf("got %q, want the post-consent retry result", got) + } +} + +// TestMCPTool_Execute_CallTimeNoToken_GateGivesUp: when Await returns an error +// (timeout / cancel / no requesting user), the call fails as no_token — no +// second CallTool, no regression from prior fail-hard behavior. +func TestMCPTool_Execute_CallTimeNoToken_GateGivesUp(t *testing.T) { + t.Parallel() + c := &mockClient{err: mcp.ErrNoToken} + gate := &gateStub{err: errors.New("consent timed out")} + a := newAdapter(t, c, func(m *MCPTool) { m.authGate = gate }) + + if _, err := a.Execute(context.Background(), json.RawMessage(`{}`)); err == nil { + t.Fatal("a gate give-up must surface as an error") + } + if gate.calls.Load() != 1 { + t.Fatalf("gate consulted %d times, want exactly 1", gate.calls.Load()) + } +} + +// TestMCPTool_Execute_CallTimeError_NoSpuriousPark: a NON-ErrNoToken CallTool +// error returns immediately without touching the gate. +func TestMCPTool_Execute_CallTimeError_NoSpuriousPark(t *testing.T) { + t.Parallel() + c := &mockClient{err: errors.New("upstream 500")} + gate := &gateStub{} + a := newAdapter(t, c, func(m *MCPTool) { m.authGate = gate }) + + if _, err := a.Execute(context.Background(), json.RawMessage(`{}`)); err == nil { + t.Fatal("a non-auth CallTool error must surface") + } + if gate.calls.Load() != 0 { + t.Fatalf("gate consulted %d times for a non-auth error, want 0", gate.calls.Load()) + } +} + +// TestMCPTool_Execute_NoGate_NoTokenSurfaces: with no gate wired, a call-time +// ErrNoToken surfaces as before (nil-gate safety). +func TestMCPTool_Execute_NoGate_NoTokenSurfaces(t *testing.T) { + t.Parallel() + a := newAdapter(t, &mockClient{err: mcp.ErrNoToken}) + if _, err := a.Execute(context.Background(), json.RawMessage(`{}`)); !errors.Is(err, mcp.ErrNoToken) { + t.Fatalf("want ErrNoToken to surface with no gate, got %v", err) + } +} From 593fa9c46f00f95e654d5602eaa113930a431a0b Mon Sep 17 00:00:00 2001 From: MK Date: Sat, 1 Aug 2026 00:41:33 -0400 Subject: [PATCH 2/3] fix(runtime): dedup only a TRAILING identical user turn on session recovery (closes #378) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The recovered-session dedup skipped the incoming user message when it matched the last user message ANYWHERE in history. A verbatim re-run of an already-answered request (the session ends in an assistant turn) therefore matched the earlier user turn and was dropped — the loop replayed the poisoned transcript and repeated its last reply without re-attempting tools. Field: a delegated MCP (Jira) call that failed "not connected" never recovered when the user connected the account and re-sent the identical request. Dedup is now trailing-only: skip only when the recovered session's LAST message is an identical user turn (a genuine premature-loop-exit duplicate — persisted but never answered). If it ends in an assistant turn, an identical incoming message is a legitimate re-run and is appended so the loop re-enters and re-attempts tools. Tests: a re-run after an answered session appends the turn (regression); a trailing-user duplicate is still skipped (no regression). --- forge-core/runtime/loop.go | 27 ++--- .../runtime/loop_session_recovery_test.go | 101 ++++++++++++++++++ 2 files changed, 116 insertions(+), 12 deletions(-) diff --git a/forge-core/runtime/loop.go b/forge-core/runtime/loop.go index e36b674..da11579 100644 --- a/forge-core/runtime/loop.go +++ b/forge-core/runtime/loop.go @@ -308,21 +308,24 @@ func (e *LLMExecutor) Execute(ctx context.Context, task *a2a.Task, msg *a2a.Mess } } - // Append the new user message, but skip if the recovered session - // already ends with an identical user message (avoids duplicates - // when users retry after a premature loop exit). + // Append the new user message, but skip a genuine duplicate: a recovered + // session whose TRAILING message is an identical user turn — the + // premature-loop-exit case where the turn was persisted but never + // answered, so a retry resends it. + // + // The dedup MUST be trailing-only. An identical user turn EARLIER in + // history that was already answered (the session ends in an assistant + // message) is NOT a duplicate — the user is legitimately re-running the + // same request (e.g. after connecting a delegated MCP account). Matching + // the last user message anywhere dropped that re-run, so the loop + // replayed the poisoned transcript and repeated the last answer without + // re-attempting tools (#378). newMsg := a2aMessageToLLM(*msg) if recovered { msgs := mem.Messages() - // Find the last user message in the recovered session. - lastUserIdx := -1 - for j := len(msgs) - 1; j >= 0; j-- { - if msgs[j].Role == llm.RoleUser { - lastUserIdx = j - break - } - } - if lastUserIdx < 0 || msgs[lastUserIdx].Content != newMsg.Content { + n := len(msgs) + trailingDup := n > 0 && msgs[n-1].Role == llm.RoleUser && msgs[n-1].Content == newMsg.Content + if !trailingDup { mem.Append(newMsg) } } else { diff --git a/forge-core/runtime/loop_session_recovery_test.go b/forge-core/runtime/loop_session_recovery_test.go index 48abb0b..13fa037 100644 --- a/forge-core/runtime/loop_session_recovery_test.go +++ b/forge-core/runtime/loop_session_recovery_test.go @@ -351,3 +351,104 @@ func strconvItoa(n int) string { } return string(buf[i:]) } + +// TestExecute_RecoveredSession_VerbatimReRunIsAppended is the #378 regression: +// a recovered session that was ALREADY ANSWERED (ends in an assistant turn) +// plus an identical incoming user message must APPEND the message and re-enter +// the loop. The old dedup matched the last user message anywhere in history, +// so a verbatim re-run (same text) was dropped — the loop replayed the +// answered transcript and repeated the last reply without re-attempting tools +// (field: a failed Jira exec never recovered after the account was connected). +func TestExecute_RecoveredSession_VerbatimReRunIsAppended(t *testing.T) { + store, err := NewMemoryStore(t.TempDir()) + if err != nil { + t.Fatalf("NewMemoryStore: %v", err) + } + const prompt = "give me top 3 open tickets under INIT Project" + // Seed an ANSWERED session: user asked, agent refused (ends in assistant). + if err := store.Save(&SessionData{ + TaskID: "rerun", Messages: []llm.ChatMessage{ + {Role: llm.RoleUser, Content: prompt}, + {Role: llm.RoleAssistant, Content: "Jira access isn't connected yet."}, + }, + }); err != nil { + t.Fatalf("Save: %v", err) + } + + var captured []llm.ChatMessage + client := &mockLLMClient{chatFunc: func(_ context.Context, req *llm.ChatRequest) (*llm.ChatResponse, error) { + if captured == nil { + captured = req.Messages + } + return &llm.ChatResponse{ + Message: llm.ChatMessage{Role: llm.RoleAssistant, Content: "here are the tickets"}, + FinishReason: "stop", + }, nil + }} + exec := NewLLMExecutor(LLMExecutorConfig{ + Client: client, MaxIterations: 5, ModelName: "test", Provider: "openai", Store: store, + }) + task := &a2a.Task{ID: "rerun"} + msg := &a2a.Message{Role: a2a.MessageRoleUser, Parts: []a2a.Part{{Kind: a2a.PartKindText, Text: prompt}}} + if _, err := exec.Execute(context.Background(), task, msg); err != nil { + t.Fatalf("Execute: %v", err) + } + + // The loop must have seen the re-run as the TRAILING turn — otherwise the + // model is prompted to continue from its own last (refusal) message. + if len(captured) == 0 { + t.Fatal("the loop never called the model on a recovered re-run") + } + last := captured[len(captured)-1] + if last.Role != llm.RoleUser || last.Content != prompt { + t.Fatalf("recovered re-run not appended: last message is %s/%q, want user/%q\nfull: %s", + last.Role, last.Content, prompt, summarizeRoles(captured)) + } +} + +// TestExecute_RecoveredSession_TrailingUserDupSkipped is the no-regression +// guard: a session ending in a USER turn (premature loop exit — persisted but +// never answered) plus an identical incoming message must NOT duplicate that +// turn (the case the dedup was written for). +func TestExecute_RecoveredSession_TrailingUserDupSkipped(t *testing.T) { + store, err := NewMemoryStore(t.TempDir()) + if err != nil { + t.Fatalf("NewMemoryStore: %v", err) + } + const prompt = "summarize the repo" + if err := store.Save(&SessionData{ + TaskID: "dup", Messages: []llm.ChatMessage{{Role: llm.RoleUser, Content: prompt}}, + }); err != nil { + t.Fatalf("Save: %v", err) + } + + var captured []llm.ChatMessage + client := &mockLLMClient{chatFunc: func(_ context.Context, req *llm.ChatRequest) (*llm.ChatResponse, error) { + if captured == nil { + captured = req.Messages + } + return &llm.ChatResponse{ + Message: llm.ChatMessage{Role: llm.RoleAssistant, Content: "done"}, + FinishReason: "stop", + }, nil + }} + exec := NewLLMExecutor(LLMExecutorConfig{ + Client: client, MaxIterations: 5, ModelName: "test", Provider: "openai", Store: store, + }) + task := &a2a.Task{ID: "dup"} + msg := &a2a.Message{Role: a2a.MessageRoleUser, Parts: []a2a.Part{{Kind: a2a.PartKindText, Text: prompt}}} + if _, err := exec.Execute(context.Background(), task, msg); err != nil { + t.Fatalf("Execute: %v", err) + } + + users := 0 + for _, m := range captured { + if m.Role == llm.RoleUser && m.Content == prompt { + users++ + } + } + if users != 1 { + t.Fatalf("trailing-user duplicate must be skipped: saw %d copies of the prompt\nfull: %s", + users, summarizeRoles(captured)) + } +} From c519976b27b5d9e24589e85e5a512d596dd9829d Mon Sep 17 00:00:00 2001 From: MK Date: Sat, 1 Aug 2026 01:03:56 -0400 Subject: [PATCH 3/3] fix(mcp): cap delegated token cache TTL so disconnect is enforced (#380) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A managed delegated (type=user) access token was cached for the platform's full expires_in (Atlassian ~1h). A platform-side disconnect deletes the vault grant but can't reach forge's in-memory cache, so the agent kept acting as the revoked user until the token lapsed (field report). Cap the per-subject delegated cache at delegatedTokenMaxTTL (5m): forge re-validates against the platform — the grant authority — within that window, so a disconnected grant yields ErrNoToken (and with the call-time auth gate, parks) in ≤5m instead of ~1h. Agent-principal (platform-mode) tokens are not capped — that grant isn't a per-user connection. Instant revocation via a platform→forge /mcp/revoke signal is the documented follow-up (#380). Test injects a spy SubjectTokenStore and asserts the 1h expires_in is clamped to the cap. --- forge-core/mcp/platform_delegated_test.go | 43 +++++++++++++++++++++++ forge-core/mcp/platform_token.go | 16 +++++++++ 2 files changed, 59 insertions(+) diff --git a/forge-core/mcp/platform_delegated_test.go b/forge-core/mcp/platform_delegated_test.go index 88d94aa..3ce09a6 100644 --- a/forge-core/mcp/platform_delegated_test.go +++ b/forge-core/mcp/platform_delegated_test.go @@ -9,6 +9,7 @@ import ( "net/http/httptest" "sync/atomic" "testing" + "time" "github.com/initializ/forge/forge-core/auth" "github.com/initializ/forge/forge-core/types" @@ -117,3 +118,45 @@ func TestBuildAuthFn_User_ResolvesSubjectFromContext(t *testing.T) { t.Fatalf("no user in ctx must fail lazily with ErrNoToken, got: %v", err) } } + +// ttlSpyStore records the ttl passed to Put, so a test can assert the cap. +type ttlSpyStore struct { + lastTTL time.Duration + toks map[string]string +} + +func (s *ttlSpyStore) Get(subject string) (string, bool) { + t, ok := s.toks[subject] + return t, ok +} +func (s *ttlSpyStore) Put(subject, token string, ttl time.Duration) { + s.lastTTL = ttl + if s.toks == nil { + s.toks = map[string]string{} + } + s.toks[subject] = token +} +func (s *ttlSpyStore) Evict(subject string) { delete(s.toks, subject) } + +// TestDelegatedTokenSource_CachesCappedTTL is the #380 fix: a platform +// disconnect can't reach forge's in-memory cache, so the delegated token is +// cached for at most delegatedTokenMaxTTL — well under the provider's ~1h +// expires_in — forcing re-validation against the platform (the grant +// authority) within that window. +func TestDelegatedTokenSource_CachesCappedTTL(t *testing.T) { + srv := delegatedServer(t, nil, "") // stub returns expires_in: 3600 (1h) + defer srv.Close() + + spy := &ttlSpyStore{} + d := newDelegatedTokenSource(PlatformSourceConfig{ + TokenEndpoint: srv.URL, AgentIdentity: "agent-cred", Ref: "atlassian", + HTTPClient: srv.Client(), SubjectStore: spy, + }) + if _, err := d.TokenForSubject(context.Background(), "alice@corp.com"); err != nil { + t.Fatalf("TokenForSubject: %v", err) + } + if spy.lastTTL != delegatedTokenMaxTTL { + t.Fatalf("cached ttl = %s, want the cap %s (provider expires_in 1h must be clamped)", + spy.lastTTL, delegatedTokenMaxTTL) + } +} diff --git a/forge-core/mcp/platform_token.go b/forge-core/mcp/platform_token.go index a2ad400..6b235af 100644 --- a/forge-core/mcp/platform_token.go +++ b/forge-core/mcp/platform_token.go @@ -33,6 +33,17 @@ const platformTokenSkew = 30 * time.Second // defaultPlatformTokenTTL applies when the endpoint omits expires_in. const defaultPlatformTokenTTL = 5 * time.Minute +// delegatedTokenMaxTTL caps how long a per-USER (delegated) access token is +// cached, independent of the provider's longer lifetime (Atlassian ≈ 1h). The +// platform is the grant authority — a disconnect deletes the grant there but +// can't reach this in-memory cache — so we re-validate against the platform at +// least this often. After a disconnect the next re-fetch (≤ this window) gets +// ErrNoToken and the call is denied/parked, instead of the agent acting as a +// revoked user for the full token lifetime (#380). Agent-principal (platform- +// mode) tokens are NOT capped here — that grant is org registration, not a +// per-user connection a user disconnects. +const delegatedTokenMaxTTL = 5 * time.Minute + type platformTokenSource struct { endpoint string // raw, ${VAR}-expandable identity string // raw, ${VAR}-expandable @@ -303,6 +314,11 @@ func (d *delegatedTokenSource) TokenForSubject(ctx context.Context, subject stri return "", fmt.Errorf("%w: platform token endpoint returned %d for server %q (subject %q)", ErrProtocolError, status, d.ref, subject) } + // Cap the cache lifetime so a platform-side disconnect is enforced within + // delegatedTokenMaxTTL, not the provider token's full lifetime (#380). + if ttl > delegatedTokenMaxTTL { + ttl = delegatedTokenMaxTTL + } d.store.Put(subject, tok, ttl) return tok, nil }