diff --git a/CHANGELOG.md b/CHANGELOG.md index d355b1e..0722e15 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,6 +25,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Scoped workflow `GITHUB_TOKEN` permissions to least privilege: read-only at the top level, per-job escalation only where needed. +### Fixed +- MCP proxy (http/sse): concurrent callers no longer share one upstream + session. The proxy now runs the MCP `initialize` handshake, sends + `Mcp-Session-Id`, keeps one session per `session_id` request param, + closes idle sessions, and recovers when the server expires one. + Streamable HTTP servers that require sessions now work through forager. + ## [0.1.0] - 2026-05-23 ### Added diff --git a/docs/proxy-modules.md b/docs/proxy-modules.md index 81601c7..5ff9e67 100644 --- a/docs/proxy-modules.md +++ b/docs/proxy-modules.md @@ -53,6 +53,14 @@ Forwards JSON-RPC requests to MCP (Model Context Protocol) servers. Supports thr **Auth types:** `basic`, `bearer`, `custom_header`, `api_key` **Creds (api_key):** `api_key_name`, `api_key_value`, `api_key_location` (`header` or `query`) +### http/sse session handling + +- Opens an MCP session with the `initialize` handshake (plus `notifications/initialized`) and sends `Mcp-Session-Id` on later requests; servers that return no session id are used sessionless +- One upstream session per `session_id` request param, so concurrent conversations don't share state on stateful servers (browser, shell, database); requests without `session_id` share one session +- Sessions close after 30 min idle (each use extends it) and on proxy close, via a best-effort HTTP `DELETE` +- A `404` for a sent session (server restarted or expired it) re-initializes and retries once +- Replies framed as SSE are unwrapped to the JSON-RPC payload + ### stdio transport details - Lazy-starts the subprocess on first request @@ -61,6 +69,7 @@ Forwards JSON-RPC requests to MCP (Model Context Protocol) servers. Supports thr - Process lifecycle: SIGTERM → 5s grace → SIGKILL - Stderr is captured and logged - Health check: `Signal(0)` to verify process is alive +- One process per datasource, shared by all callers: `session_id` does not isolate state on stdio ## MongoDB Proxy (`mongo-proxy`) diff --git a/pkg/proxy/mcp/proxy.go b/pkg/proxy/mcp/proxy.go index b6cc5a7..e098432 100644 --- a/pkg/proxy/mcp/proxy.go +++ b/pkg/proxy/mcp/proxy.go @@ -31,6 +31,14 @@ type Proxy struct { workingDir string stdio *stdioProcess + // http/sse upstream sessions, keyed by the caller's session_id; sessionMu + // guards every field in this group. + sessionMu sync.Mutex + sessions map[string]*mcpSession + sessionlessUntil time.Time + stopSweep chan struct{} + closed bool + // oauth2 token cache oauthMu sync.Mutex oauthToken string @@ -43,7 +51,8 @@ type Proxy struct { // New creates a new MCP proxy. func New(logger *slog.Logger) *Proxy { return &Proxy{ - logger: logger, + logger: logger, + sessions: make(map[string]*mcpSession), client: &http.Client{ Timeout: 120 * time.Second, }, @@ -74,6 +83,12 @@ func (p *Proxy) Configure(config map[string]any, creds map[string]string) error if p.url == "" { return fmt.Errorf("url is required for MCP %s transport", p.transport) } + p.sessionMu.Lock() + if p.stopSweep == nil && !p.closed { + p.stopSweep = make(chan struct{}) + go p.sweepSessions(p.stopSweep) + } + p.sessionMu.Unlock() case "stdio": if v, ok := config["command"].(string); ok { p.command = v @@ -104,10 +119,8 @@ func (p *Proxy) Configure(config map[string]any, creds map[string]string) error func (p *Proxy) HandleRequest(ctx context.Context, req *proxy.ActionRequest) (*proxy.ActionResponse, error) { switch p.transport { - case "http": - return p.handleHTTP(ctx, req) - case "sse": - return p.handleSSE(ctx, req) + case "http", "sse": + return p.handleWithSession(ctx, req) case "stdio": return p.handleStdio(ctx, req) default: @@ -138,67 +151,9 @@ func (p *Proxy) buildRequestBody(req *proxy.ActionRequest) ([]byte, error) { return nil, fmt.Errorf("no request body or params provided") } -func (p *Proxy) handleHTTP(ctx context.Context, req *proxy.ActionRequest) (*proxy.ActionResponse, error) { - body, err := p.buildRequestBody(req) - if err != nil { - return nil, fmt.Errorf("building MCP request body: %w", err) - } - - httpReq, err := http.NewRequestWithContext(ctx, "POST", p.url, bytes.NewReader(body)) - if err != nil { - return nil, fmt.Errorf("creating MCP request: %w", err) - } - httpReq.Header.Set("Content-Type", "application/json") - if err := p.injectAuth(httpReq); err != nil { - return nil, fmt.Errorf("MCP auth: %w", err) - } - - resp, err := p.client.Do(httpReq) - if err != nil { - return nil, fmt.Errorf("MCP request failed: %w", err) - } - defer func() { _ = resp.Body.Close() }() - - respBody, err := io.ReadAll(resp.Body) - if err != nil { - return nil, fmt.Errorf("reading MCP response: %w", err) - } - - return &proxy.ActionResponse{ - StatusCode: resp.StatusCode, - Data: string(respBody), - }, nil -} - -func (p *Proxy) handleSSE(ctx context.Context, req *proxy.ActionRequest) (*proxy.ActionResponse, error) { - body, err := p.buildRequestBody(req) - if err != nil { - return nil, fmt.Errorf("building MCP request body: %w", err) - } - - httpReq, err := http.NewRequestWithContext(ctx, "POST", p.url, bytes.NewReader(body)) - if err != nil { - return nil, fmt.Errorf("creating MCP SSE request: %w", err) - } - httpReq.Header.Set("Content-Type", "application/json") - httpReq.Header.Set("Accept", "text/event-stream") - if err := p.injectAuth(httpReq); err != nil { - return nil, fmt.Errorf("MCP SSE auth: %w", err) - } - - resp, err := p.client.Do(httpReq) - if err != nil { - return nil, fmt.Errorf("MCP SSE request failed: %w", err) - } - defer func() { _ = resp.Body.Close() }() - - // Read entire body first to support fallback for non-SSE responses - respBody, err := io.ReadAll(resp.Body) - if err != nil { - return nil, fmt.Errorf("reading MCP SSE response: %w", err) - } - - // Parse SSE response — collect all "data:" lines into a single JSON-RPC response +// parseSSE collects the "data:" lines of an SSE-framed response into one +// JSON-RPC payload, falling back to the raw body when there are none. +func parseSSE(respBody []byte) (string, error) { var result strings.Builder scanner := bufio.NewScanner(bytes.NewReader(respBody)) for scanner.Scan() { @@ -212,19 +167,12 @@ func (p *Proxy) handleSSE(ctx context.Context, req *proxy.ActionRequest) (*proxy } } if err := scanner.Err(); err != nil { - return nil, fmt.Errorf("reading MCP SSE stream: %w", err) + return "", fmt.Errorf("reading MCP SSE stream: %w", err) } - - responseData := result.String() - if responseData == "" { - // Fallback: use entire response body if no SSE data lines were found - responseData = string(respBody) + if result.Len() == 0 { + return string(respBody), nil } - - return &proxy.ActionResponse{ - StatusCode: resp.StatusCode, - Data: responseData, - }, nil + return result.String(), nil } func (p *Proxy) handleStdio(ctx context.Context, req *proxy.ActionRequest) (*proxy.ActionResponse, error) { @@ -399,6 +347,23 @@ func (p *Proxy) HealthCheck(ctx context.Context) error { } func (p *Proxy) Close() error { + p.sessionMu.Lock() + p.closed = true + if p.stopSweep != nil { + close(p.stopSweep) + p.stopSweep = nil + } + p.sessionMu.Unlock() + // In parallel so a hung MCP server delays shutdown by one timeout, not one per session. + var wg sync.WaitGroup + for _, id := range p.takeSessions(func(*mcpSession) bool { return true }) { + wg.Add(1) + go func() { + defer wg.Done() + p.terminateSession(id) + }() + } + wg.Wait() p.client.CloseIdleConnections() if p.stdio != nil { return p.stdio.close() diff --git a/pkg/proxy/mcp/session.go b/pkg/proxy/mcp/session.go new file mode 100644 index 0000000..4c6c4a8 --- /dev/null +++ b/pkg/proxy/mcp/session.go @@ -0,0 +1,299 @@ +package mcp + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "strings" + "time" + + "nudgebee/forager/pkg/proxy" +) + +// sessionIdleTTL is how long an upstream MCP session may sit unused before it is +// closed. Each use pushes the deadline out again, so an active conversation +// keeps its session. +const sessionIdleTTL = 30 * time.Minute + +const ( + sessionSweepInterval = 5 * time.Minute + sessionTerminateTimeout = 10 * time.Second +) + +var errProxyClosed = errors.New("MCP proxy is closed") + +type mcpSession struct { + id string + lastUsed time.Time +} + +// sessionKey scopes the upstream session to the caller's conversation. Callers +// that send no session_id share one session, as before. +func sessionKey(req *proxy.ActionRequest) string { + if req.Params == nil { + return "" + } + key, _ := req.Params["session_id"].(string) + return key +} + +// handleWithSession runs an HTTP/SSE request inside the upstream MCP session for +// the caller's session_id, opening one with the initialize handshake when +// needed. Stateful MCP servers (browser, shell, database session) keep state per +// session, so without this every conversation would act on the same state. +func (p *Proxy) handleWithSession(ctx context.Context, req *proxy.ActionRequest) (*proxy.ActionResponse, error) { + body, err := p.buildRequestBody(req) + if err != nil { + return nil, fmt.Errorf("building MCP request body: %w", err) + } + + key := sessionKey(req) + sessionID, err := p.ensureSession(ctx, key) + if errors.Is(err, errProxyClosed) { + return nil, err + } + if err != nil { + // Servers that don't implement sessions may reject initialize; the + // request itself can still succeed without one. + p.logger.Warn("MCP session init failed, proceeding without session", "error", err) + } + + resp, err := p.post(ctx, body, sessionID) + if err != nil { + return nil, err + } + if !sessionGone(resp.ActionResponse, sessionID) { + return resp.ActionResponse, nil + } + + p.logger.Info("MCP session expired upstream, re-initializing", "status", resp.StatusCode) + p.dropSession(key, sessionID) + sessionID, err = p.ensureSession(ctx, key) + if err != nil { + return nil, fmt.Errorf("MCP session re-init failed: %w", err) + } + resp, err = p.post(ctx, body, sessionID) + if err != nil { + return nil, err + } + return resp.ActionResponse, nil +} + +// sessionGone reports whether the server no longer knows the session we sent. +// Only 404, as the MCP spec prescribes: guessing from a 400's text would drop a +// live session and replay a tool call on an ordinary tool error. +func sessionGone(resp *proxy.ActionResponse, sessionID string) bool { + return sessionID != "" && resp.StatusCode == http.StatusNotFound +} + +// ensureSession returns the live session for key, extending its idle deadline, +// or opens a new one. It returns "" when the server does not use sessions. +func (p *Proxy) ensureSession(ctx context.Context, key string) (string, error) { + p.sessionMu.Lock() + if p.closed { + p.sessionMu.Unlock() + return "", errProxyClosed + } + if time.Now().Before(p.sessionlessUntil) { + p.sessionMu.Unlock() + return "", nil + } + if s, ok := p.sessions[key]; ok { + if time.Since(s.lastUsed) < sessionIdleTTL { + s.lastUsed = time.Now() + p.sessionMu.Unlock() + return s.id, nil + } + delete(p.sessions, key) + go p.terminateSession(s.id) + } + p.sessionMu.Unlock() + + // Handshake outside the lock so one slow server doesn't block other sessions. + id, sessionless, err := p.initializeSession(ctx) + if sessionless { + // Remember it, or every request would pay an extra initialize round-trip. + p.sessionMu.Lock() + p.sessionlessUntil = time.Now().Add(sessionIdleTTL) + p.sessionMu.Unlock() + } + if err != nil || id == "" { + return "", err + } + + p.sessionMu.Lock() + if p.closed { + p.sessionMu.Unlock() + go p.terminateSession(id) + return "", errProxyClosed + } + if s, ok := p.sessions[key]; ok && time.Since(s.lastUsed) < sessionIdleTTL { + // A concurrent request for the same key won the race: share its + // session and close the one we just opened. + s.lastUsed = time.Now() + p.sessionMu.Unlock() + go p.terminateSession(id) + return s.id, nil + } + p.sessions[key] = &mcpSession{id: id, lastUsed: time.Now()} + p.sessionMu.Unlock() + return id, nil +} + +// dropSession forgets key only if it still maps to sessionID, so a session a +// concurrent request has just opened is not discarded. +func (p *Proxy) dropSession(key, sessionID string) { + p.sessionMu.Lock() + if s, ok := p.sessions[key]; ok && s.id == sessionID { + delete(p.sessions, key) + } + p.sessionMu.Unlock() +} + +// initializeSession performs the MCP initialize handshake and returns the +// server-assigned Mcp-Session-Id. sessionless is true only when the server +// answered definitively without one (success with no id, or a 4xx refusing +// initialize); transport errors and 5xx are not, so a blip can't disable +// sessions for a server that needs them. +func (p *Proxy) initializeSession(ctx context.Context) (id string, sessionless bool, err error) { + initBody, _ := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": 0, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": "2025-03-26", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "forager", "version": "1.0.0"}, + }, + }) + resp, err := p.post(ctx, initBody, "") + if err != nil { + return "", false, fmt.Errorf("initialize: %w", err) + } + if resp.StatusCode >= 400 { + return "", resp.StatusCode < 500, fmt.Errorf("initialize returned HTTP %d", resp.StatusCode) + } + id = resp.sessionID + if id == "" { + return "", true, nil + } + + // The spec requires this notification before other requests; servers that + // don't enforce it ignore it, so a failure here is not fatal. + notify := []byte(`{"jsonrpc":"2.0","method":"notifications/initialized"}`) + if _, err := p.post(ctx, notify, id); err != nil { + p.logger.Debug("MCP initialized notification failed", "error", err) + } + return id, false, nil +} + +// terminateSession asks the server to end a session (HTTP DELETE with +// Mcp-Session-Id). Best-effort: servers may answer 404 (already gone) or 405 +// (client termination unsupported); the server's own timeout is the backstop. +func (p *Proxy) terminateSession(sessionID string) { + ctx, cancel := context.WithTimeout(context.Background(), sessionTerminateTimeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodDelete, p.url, nil) + if err != nil { + return + } + req.Header.Set("Mcp-Session-Id", sessionID) + if err := p.injectAuth(req); err != nil { + p.logger.Debug("MCP session terminate: auth failed", "error", err) + return + } + resp, err := p.client.Do(req) + if err != nil { + p.logger.Debug("MCP session terminate failed", "error", err) + return + } + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() +} + +// sweepSessions closes sessions idle past sessionIdleTTL until stop is closed. +// Terminations run one at a time, which bounds load on the MCP server. +func (p *Proxy) sweepSessions(stop <-chan struct{}) { + ticker := time.NewTicker(sessionSweepInterval) + defer ticker.Stop() + for { + select { + case <-stop: + return + case <-ticker.C: + for _, id := range p.takeSessions(func(s *mcpSession) bool { + return time.Since(s.lastUsed) >= sessionIdleTTL + }) { + p.terminateSession(id) + } + } + } +} + +// takeSessions removes and returns the ids of sessions matching match. +func (p *Proxy) takeSessions(match func(*mcpSession) bool) []string { + p.sessionMu.Lock() + defer p.sessionMu.Unlock() + var ids []string + for key, s := range p.sessions { + if match(s) { + delete(p.sessions, key) + ids = append(ids, s.id) + } + } + return ids +} + +// httpResponse is an ActionResponse plus the Mcp-Session-Id header, which the +// initialize handshake needs. +type httpResponse struct { + *proxy.ActionResponse + sessionID string +} + +// post sends one JSON-RPC message, with the session header when sessionID is +// set, and unwraps an SSE-framed reply into the JSON-RPC payload. +func (p *Proxy) post(ctx context.Context, body []byte, sessionID string) (httpResponse, error) { + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, p.url, bytes.NewReader(body)) + if err != nil { + return httpResponse{}, fmt.Errorf("creating MCP request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + // Streamable HTTP servers reject requests that don't accept both. + httpReq.Header.Set("Accept", "application/json, text/event-stream") + if sessionID != "" { + httpReq.Header.Set("Mcp-Session-Id", sessionID) + } + if err := p.injectAuth(httpReq); err != nil { + return httpResponse{}, fmt.Errorf("MCP auth: %w", err) + } + + resp, err := p.client.Do(httpReq) + if err != nil { + return httpResponse{}, fmt.Errorf("MCP request failed: %w", err) + } + defer func() { _ = resp.Body.Close() }() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return httpResponse{}, fmt.Errorf("reading MCP response: %w", err) + } + + data := string(respBody) + if p.transport == "sse" || strings.Contains(resp.Header.Get("Content-Type"), "text/event-stream") { + data, err = parseSSE(respBody) + if err != nil { + return httpResponse{}, err + } + } + + return httpResponse{ + ActionResponse: &proxy.ActionResponse{StatusCode: resp.StatusCode, Data: data}, + sessionID: resp.Header.Get("Mcp-Session-Id"), + }, nil +} diff --git a/pkg/proxy/mcp/session_test.go b/pkg/proxy/mcp/session_test.go new file mode 100644 index 0000000..913337d --- /dev/null +++ b/pkg/proxy/mcp/session_test.go @@ -0,0 +1,245 @@ +package mcp + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "os" + "sync" + "testing" + "time" + + "nudgebee/forager/pkg/proxy" +) + +// fakeServer is a minimal Streamable HTTP MCP server: initialize issues a new +// Mcp-Session-Id, requests with an unknown session get 404 (per the MCP spec), +// requests that don't accept both JSON and SSE get 406, and DELETE ends a session. +type fakeServer struct { + mu sync.Mutex + next int + live map[string]bool + notified map[string]bool + deleted []string + // sessionless makes initialize succeed without issuing a session id. + sessionless bool + inits int +} + +func (f *fakeServer) ServeHTTP(w http.ResponseWriter, r *http.Request) { + f.mu.Lock() + defer f.mu.Unlock() + sid := r.Header.Get("Mcp-Session-Id") + if r.Method == http.MethodDelete { + f.deleted = append(f.deleted, sid) + delete(f.live, sid) + return + } + if r.Header.Get("Accept") != "application/json, text/event-stream" { + w.WriteHeader(http.StatusNotAcceptable) + return + } + body, _ := io.ReadAll(r.Body) + var rpc struct { + Method string `json:"method"` + } + _ = json.Unmarshal(body, &rpc) + switch { + case rpc.Method == "initialize" && f.sessionless: + f.inits++ + _, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":0,"result":{}}`)) + case f.sessionless: + _, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":1,"result":{}}`)) + case rpc.Method == "initialize": + f.inits++ + f.next++ + sid = fmt.Sprintf("s%d", f.next) + f.live[sid] = true + w.Header().Set("Mcp-Session-Id", sid) + _, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":0,"result":{}}`)) + case !f.live[sid]: + w.WriteHeader(http.StatusNotFound) + case rpc.Method == "notifications/initialized": + f.notified[sid] = true + w.WriteHeader(http.StatusAccepted) + default: + // Reply SSE-framed, as Streamable HTTP servers may. + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprintf(w, "event: message\ndata: {\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{\"session\":%q}}\n\n", sid) + } +} + +func (f *fakeServer) deletedIDs() []string { + f.mu.Lock() + defer f.mu.Unlock() + return append([]string(nil), f.deleted...) +} + +func newTestProxy(t *testing.T) (*Proxy, *fakeServer) { + t.Helper() + fake := &fakeServer{live: map[string]bool{}, notified: map[string]bool{}} + srv := httptest.NewServer(fake) + p := New(slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError}))) + if err := p.Configure(map[string]any{"transport": "http", "url": srv.URL}, nil); err != nil { + t.Fatalf("Configure: %v", err) + } + t.Cleanup(func() { + _ = p.Close() + srv.Close() + }) + return p, fake +} + +func toolCall(t *testing.T, p *Proxy, sessionID string) string { + t.Helper() + params := map[string]any{"method": "tools/call", "params": map[string]any{"name": "x"}} + if sessionID != "" { + params["session_id"] = sessionID + } + resp, err := p.HandleRequest(context.Background(), &proxy.ActionRequest{Action: "mcp_request", Params: params}) + if err != nil { + t.Fatalf("HandleRequest: %v", err) + } + if resp.StatusCode != http.StatusOK { + t.Fatalf("status %d: %s", resp.StatusCode, resp.Data) + } + return resp.Data +} + +func result(sid string) string { + return fmt.Sprintf(`{"jsonrpc":"2.0","id":1,"result":{"session":%q}}`, sid) +} + +func TestSession_EachCallerSessionGetsItsOwnUpstreamSession(t *testing.T) { + p, fake := newTestProxy(t) + + a1 := toolCall(t, p, "conv-a") + b := toolCall(t, p, "conv-b") + a2 := toolCall(t, p, "conv-a") + + if a1 != result("s1") || a2 != result("s1") { + t.Fatalf("conv-a should keep session s1, got %s then %s", a1, a2) + } + if b != result("s2") { + t.Fatalf("conv-b should get its own session s2, got %s", b) + } + fake.mu.Lock() + defer fake.mu.Unlock() + if !fake.notified["s1"] || !fake.notified["s2"] { + t.Fatalf("initialized notification not sent: %v", fake.notified) + } +} + +func TestSession_NoSessionIDSharesOneSession(t *testing.T) { + p, _ := newTestProxy(t) + if first, second := toolCall(t, p, ""), toolCall(t, p, ""); first != second { + t.Fatalf("calls without session_id should share a session: %s vs %s", first, second) + } +} + +func TestSession_Upstream404ReinitializesAndRetries(t *testing.T) { + p, fake := newTestProxy(t) + toolCall(t, p, "conv") + + fake.mu.Lock() + fake.live = map[string]bool{} // server restarted, forgot every session + fake.mu.Unlock() + + if got := toolCall(t, p, "conv"); got != result("s2") { + t.Fatalf("expected retry on fresh session s2, got %s", got) + } +} + +func TestSession_IdleSessionIsTerminatedAndReplaced(t *testing.T) { + p, fake := newTestProxy(t) + toolCall(t, p, "conv") + + p.sessionMu.Lock() + p.sessions["conv"].lastUsed = time.Now().Add(-sessionIdleTTL - time.Second) + p.sessionMu.Unlock() + + if got := toolCall(t, p, "conv"); got != result("s2") { + t.Fatalf("idle session should be replaced, got %s", got) + } + deadline := time.Now().Add(2 * time.Second) + for len(fake.deletedIDs()) == 0 && time.Now().Before(deadline) { + time.Sleep(10 * time.Millisecond) + } + if d := fake.deletedIDs(); len(d) != 1 || d[0] != "s1" { + t.Fatalf("expected DELETE of s1, got %v", d) + } +} + +func TestSession_CloseTerminatesLiveSessions(t *testing.T) { + p, fake := newTestProxy(t) + toolCall(t, p, "conv-a") + toolCall(t, p, "conv-b") + + if err := p.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + if d := fake.deletedIDs(); len(d) != 2 { + t.Fatalf("Close should DELETE both sessions, got %v", d) + } +} + +func TestSessionGone(t *testing.T) { + cases := []struct { + name string + status int + data string + sid string + want bool + }{ + {"404 with session", 404, "", "s1", true}, + {"400 mentioning session is not treated as gone", 400, "browser session not found", "s1", false}, + {"404 without session", 404, "", "", false}, + {"200", 200, "session", "s1", false}, + } + for _, c := range cases { + got := sessionGone(&proxy.ActionResponse{StatusCode: c.status, Data: c.data}, c.sid) + if got != c.want { + t.Errorf("%s: got %v, want %v", c.name, got, c.want) + } + } +} + +func TestSession_SessionlessServerIsInitializedOnce(t *testing.T) { + p, fake := newTestProxy(t) + fake.mu.Lock() + fake.sessionless = true + fake.mu.Unlock() + + for range 3 { + toolCall(t, p, "conv") + } + fake.mu.Lock() + defer fake.mu.Unlock() + if fake.inits != 1 { + t.Fatalf("sessionless server should be initialized once, got %d", fake.inits) + } +} + +func TestSession_ClosedProxyOpensNoSession(t *testing.T) { + p, fake := newTestProxy(t) + if err := p.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + if _, err := p.ensureSession(context.Background(), "conv"); err != errProxyClosed { + t.Fatalf("expected errProxyClosed, got %v", err) + } + req := &proxy.ActionRequest{Params: map[string]any{"method": "tools/call", "session_id": "conv"}} + if _, err := p.HandleRequest(context.Background(), req); !errors.Is(err, errProxyClosed) { + t.Fatalf("request on a closed proxy should fail with errProxyClosed, got %v", err) + } + fake.mu.Lock() + defer fake.mu.Unlock() + if fake.inits != 0 { + t.Fatalf("closed proxy should not initialize, got %d", fake.inits) + } +}