From 5e2504b2a36cc28a6e126432e328944b0f47cc79 Mon Sep 17 00:00:00 2001 From: Utkarsh Katiyar Date: Fri, 7 Aug 2026 18:08:48 +0530 Subject: [PATCH 1/2] feat(session): implement session persistence and status retrieval for chat sessions --- forge-core/runtime/loop.go | 5 + .../runtime/loop_refresh_persistence_test.go | 182 ++++++++++++++++++ forge-ui/chat.go | 100 ++++++++-- forge-ui/server.go | 1 + forge-ui/static/app.js | 79 +++++++- 5 files changed, 347 insertions(+), 20 deletions(-) create mode 100644 forge-core/runtime/loop_refresh_persistence_test.go diff --git a/forge-core/runtime/loop.go b/forge-core/runtime/loop.go index 6f74d84f..7fecdfaa 100644 --- a/forge-core/runtime/loop.go +++ b/forge-core/runtime/loop.go @@ -337,6 +337,8 @@ func (e *LLMExecutor) Execute(ctx context.Context, task *a2a.Task, msg *a2a.Mess mem.Append(newMsg) } + e.persistSession(task.ID, mem) + // Build tool definitions var toolDefs []llm.ToolDefinition if e.tools != nil { @@ -382,6 +384,7 @@ func (e *LLMExecutor) Execute(ctx context.Context, task *a2a.Task, msg *a2a.Mess // runner, which maps it to TaskStateCanceled + // invocation_cancelled audit. See issue #88 / FWS-4. if err := ctx.Err(); err != nil { + e.persistSession(task.ID, mem) return nil, err } @@ -457,6 +460,7 @@ func (e *LLMExecutor) Execute(ctx context.Context, task *a2a.Task, msg *a2a.Mess // route to invocation_cancelled instead of state=failed. // See issue #88 / FWS-4. if cerr := ctx.Err(); cerr != nil { + e.persistSession(task.ID, mem) return nil, cerr } _ = e.hooks.Fire(ctx, OnError, &HookContext{ @@ -702,6 +706,7 @@ func (e *LLMExecutor) Execute(ctx context.Context, task *a2a.Task, msg *a2a.Mess // orchestrators that cancel mid-iteration get fast exit // without burning more LLM/tool spend. See issue #88 / FWS-4. if err := ctx.Err(); err != nil { + e.persistSession(task.ID, mem) return nil, err } toolsUsed = append(toolsUsed, tc.Function.Name) diff --git a/forge-core/runtime/loop_refresh_persistence_test.go b/forge-core/runtime/loop_refresh_persistence_test.go new file mode 100644 index 00000000..0f1f3a81 --- /dev/null +++ b/forge-core/runtime/loop_refresh_persistence_test.go @@ -0,0 +1,182 @@ +package runtime + +import ( + "context" + "errors" + "path/filepath" + "testing" + + "github.com/initializ/forge/forge-core/a2a" + "github.com/initializ/forge/forge-core/llm" +) + +// Regression test for the "browser refresh mid-response" bug: a client +// disconnect (e.g. a page reload) cancels the context passed to Execute +// before the LLM call ever completes. Before the fix, persistSession was +// only called on the successful-completion paths, so the user's own +// message — and any tool-call progress already made — was never written +// to disk if that cancellation raced ahead of completion. This test +// drives Execute with a context that's cancelled synchronously from +// inside the (mocked) LLM call, using a REAL file-backed MemoryStore, and +// asserts the persisted session file on disk contains the user's message +// even though Execute returns context.Canceled. +func TestLLMExecutor_RefreshMidResponse_UserMessagePersisted(t *testing.T) { + store, err := NewMemoryStore(filepath.Join(t.TempDir(), "sessions")) + if err != nil { + t.Fatalf("NewMemoryStore: %v", err) + } + + var cancel context.CancelFunc + chatCalls := 0 + + exec := NewLLMExecutor(LLMExecutorConfig{ + Client: &mockLLMClient{ + chatFunc: func(_ context.Context, _ *llm.ChatRequest) (*llm.ChatResponse, error) { + chatCalls++ + // Simulate the browser tab reloading right as the agent + // starts processing the turn: the client disconnects, + // forge-ui's proxy context cancels, which (pre-fix) would + // propagate all the way here before any persist happened. + cancel() + return nil, context.Canceled + }, + }, + Tools: &mockToolExecutor{toolDefs: []llm.ToolDefinition{}}, + Store: store, + ModelName: "test", + Provider: "test", + }) + + var ctx context.Context + ctx, cancel = context.WithCancel(context.Background()) + defer cancel() + + const taskID = "refresh-test-task" + _, err = exec.Execute(ctx, + &a2a.Task{ID: taskID}, + &a2a.Message{Role: a2a.MessageRoleUser, Parts: []a2a.Part{a2a.NewTextPart("hello, are you there?")}}, + ) + if !errors.Is(err, context.Canceled) { + t.Fatalf("Execute should return context.Canceled (simulating the refresh-aborted turn), got %v", err) + } + if chatCalls != 1 { + t.Fatalf("expected exactly 1 LLM call before cancellation, got %d", chatCalls) + } + + // The critical assertion: even though the turn never completed, the + // user's message must already be on disk — durability of input must + // not depend on the response finishing. + saved, loadErr := store.Load(taskID) + if loadErr != nil { + t.Fatalf("store.Load: %v", loadErr) + } + if saved == nil { + t.Fatal("session was never persisted — the user's message is lost on refresh, exactly the bug this test guards against") + } + + found := false + for _, m := range saved.Messages { + if m.Role == llm.RoleUser && m.Content == "hello, are you there?" { + found = true + break + } + } + if !found { + t.Fatalf("persisted session does not contain the user's message; got messages: %+v", saved.Messages) + } +} + +// Companion test: once the turn eventually DOES complete (the agent kept +// working in the background per the detached-context fix in +// forge-ui/chat.go, even though the original HTTP client vanished), a +// second call to Execute against the SAME task ID must not need the +// original message again — it should already be recoverable from disk, +// and the session file must accumulate rather than reset. +func TestLLMExecutor_RefreshThenRetry_SessionAccumulates(t *testing.T) { + store, err := NewMemoryStore(filepath.Join(t.TempDir(), "sessions")) + if err != nil { + t.Fatalf("NewMemoryStore: %v", err) + } + + const taskID = "refresh-retry-task" + + // First turn: cancelled immediately (the "refresh"). + var cancel context.CancelFunc + cancelExec := NewLLMExecutor(LLMExecutorConfig{ + Client: &mockLLMClient{ + chatFunc: func(_ context.Context, _ *llm.ChatRequest) (*llm.ChatResponse, error) { + cancel() + return nil, context.Canceled + }, + }, + Tools: &mockToolExecutor{toolDefs: []llm.ToolDefinition{}}, + Store: store, + ModelName: "test", + Provider: "test", + }) + var ctx context.Context + ctx, cancel = context.WithCancel(context.Background()) + _, err = cancelExec.Execute(ctx, + &a2a.Task{ID: taskID}, + &a2a.Message{Role: a2a.MessageRoleUser, Parts: []a2a.Part{a2a.NewTextPart("first message")}}, + ) + if !errors.Is(err, context.Canceled) { + t.Fatalf("first Execute should be cancelled, got %v", err) + } + + // Second turn: a fresh, un-cancelled context (simulating the agent + // process having kept running, or the user reconnecting and the + // session recovering from disk) completes normally. + okExec := NewLLMExecutor(LLMExecutorConfig{ + Client: &mockLLMClient{ + chatFunc: func(_ context.Context, req *llm.ChatRequest) (*llm.ChatResponse, error) { + return &llm.ChatResponse{ + ID: "r2", + Message: llm.ChatMessage{Role: llm.RoleAssistant, Content: "all done"}, + }, nil + }, + }, + Tools: &mockToolExecutor{toolDefs: []llm.ToolDefinition{}}, + Store: store, + ModelName: "test", + Provider: "test", + }) + + _, err = okExec.Execute(context.Background(), + &a2a.Task{ID: taskID, History: nil}, + &a2a.Message{Role: a2a.MessageRoleUser, Parts: []a2a.Part{a2a.NewTextPart("second message")}}, + ) + if err != nil { + t.Fatalf("second Execute should succeed, got %v", err) + } + + saved, loadErr := store.Load(taskID) + if loadErr != nil { + t.Fatalf("store.Load: %v", loadErr) + } + if saved == nil { + t.Fatal("session missing after second turn") + } + + var userMsgs []string + for _, m := range saved.Messages { + if m.Role == llm.RoleUser { + userMsgs = append(userMsgs, m.Content) + } + } + if len(userMsgs) != 2 { + t.Fatalf("expected both turns' user messages preserved in one session file (accumulated, not overwritten-and-lost), got %v", userMsgs) + } + if userMsgs[0] != "first message" || userMsgs[1] != "second message" { + t.Fatalf("unexpected message order/content: %v", userMsgs) + } + + // Same task ID -> same file on disk, not a forked/orphaned session. + ids, listErr := store.List() + if listErr != nil { + t.Fatalf("store.List: %v", listErr) + } + if len(ids) != 1 { + t.Fatalf("expected exactly one session file on disk for this task ID, got %v", ids) + } +} diff --git a/forge-ui/chat.go b/forge-ui/chat.go index f27751d5..84cc712d 100644 --- a/forge-ui/chat.go +++ b/forge-ui/chat.go @@ -3,6 +3,7 @@ package forgeui import ( "bufio" "bytes" + "context" "encoding/json" "fmt" "io" @@ -16,6 +17,8 @@ import ( "github.com/initializ/forge/forge-core/auth" ) +const agentCallTimeout = 10 * time.Minute + // handleChat proxies a chat message to a running agent via A2A JSON-RPC // and streams the SSE response back to the browser. func (s *UIServer) handleChat(w http.ResponseWriter, r *http.Request) { @@ -73,10 +76,12 @@ func (s *UIServer) handleChat(w http.ResponseWriter, r *http.Request) { return } - // POST to the agent's A2A endpoint. + agentCtx, cancel := context.WithTimeout(context.WithoutCancel(r.Context()), agentCallTimeout) + defer cancel() + client := &http.Client{Timeout: 0} agentURL := fmt.Sprintf("http://127.0.0.1:%d/", port) - agentReq, err := http.NewRequestWithContext(r.Context(), http.MethodPost, agentURL, bytes.NewReader(rpcBody)) + agentReq, err := http.NewRequestWithContext(agentCtx, http.MethodPost, agentURL, bytes.NewReader(rpcBody)) if err != nil { writeError(w, http.StatusInternalServerError, "failed to create agent request") return @@ -115,17 +120,18 @@ func (s *UIServer) handleChat(w http.ResponseWriter, r *http.Request) { w.Header().Set("Connection", "keep-alive") flusher.Flush() - // Parse agent SSE and re-emit to browser. scanner := bufio.NewScanner(agentResp.Body) var eventType string var dataLines []string + clientGone := false for scanner.Scan() { - // Check if client disconnected. - select { - case <-r.Context().Done(): - return - default: + if !clientGone { + select { + case <-r.Context().Done(): + clientGone = true + default: + } } line := scanner.Text() @@ -135,11 +141,13 @@ func (s *UIServer) handleChat(w http.ResponseWriter, r *http.Request) { } else if after, found := strings.CutPrefix(line, "data:"); found { dataLines = append(dataLines, after) } else if line == "" && eventType != "" { - // Blank line = end of SSE frame. Re-emit to browser. - data := strings.TrimSpace(strings.Join(dataLines, "\n")) - if data != "" { - _, _ = fmt.Fprintf(w, "event: %s\ndata: %s\n\n", eventType, data) - flusher.Flush() + // Blank line = end of SSE frame. Re-emit to browser (if still connected). + if !clientGone { + data := strings.TrimSpace(strings.Join(dataLines, "\n")) + if data != "" { + _, _ = fmt.Fprintf(w, "event: %s\ndata: %s\n\n", eventType, data) + flusher.Flush() + } } eventType = "" dataLines = nil @@ -152,6 +160,72 @@ func (s *UIServer) handleChat(w http.ResponseWriter, r *http.Request) { flusher.Flush() } +func (s *UIServer) handleSessionStatus(w http.ResponseWriter, r *http.Request) { + agentID := r.PathValue("id") + sid := r.PathValue("sid") + if agentID == "" || sid == "" { + writeError(w, http.StatusBadRequest, "agent id and session id are required") + return + } + + agents, err := s.scanner.Scan() + if err != nil { + writeError(w, http.StatusInternalServerError, err.Error()) + return + } + agent := agents[agentID] + if agent == nil || agent.Port == 0 { + writeJSON(w, http.StatusOK, map[string]string{"state": "unknown"}) + return + } + + rpcBody, err := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "tasks/get", + "params": map[string]any{"id": sid}, + }) + if err != nil { + writeError(w, http.StatusInternalServerError, "failed to build request") + return + } + + ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second) + defer cancel() + + agentURL := fmt.Sprintf("http://127.0.0.1:%d/", agent.Port) + agentReq, err := http.NewRequestWithContext(ctx, http.MethodPost, agentURL, bytes.NewReader(rpcBody)) + if err != nil { + writeError(w, http.StatusInternalServerError, "failed to create agent request") + return + } + agentReq.Header.Set("Content-Type", "application/json") + if token := s.loadAgentToken(agentID); token != "" { + agentReq.Header.Set("Authorization", "Bearer "+token) + } + + client := &http.Client{Timeout: 6 * time.Second} + agentResp, err := client.Do(agentReq) + if err != nil { + writeJSON(w, http.StatusOK, map[string]string{"state": "unknown"}) + return + } + defer func() { _ = agentResp.Body.Close() }() + + var rpcResp struct { + Result *struct { + Status struct { + State string `json:"state"` + } `json:"status"` + } `json:"result"` + } + if err := json.NewDecoder(agentResp.Body).Decode(&rpcResp); err != nil || rpcResp.Result == nil { + writeJSON(w, http.StatusOK, map[string]string{"state": "unknown"}) + return + } + writeJSON(w, http.StatusOK, map[string]string{"state": rpcResp.Result.Status.State}) +} + // handleListSessions returns stored chat sessions for an agent. func (s *UIServer) handleListSessions(w http.ResponseWriter, r *http.Request) { agentID := r.PathValue("id") diff --git a/forge-ui/server.go b/forge-ui/server.go index 9c43684e..80284387 100644 --- a/forge-ui/server.go +++ b/forge-ui/server.go @@ -80,6 +80,7 @@ func (s *UIServer) Start(ctx context.Context) error { mux.HandleFunc("POST /api/agents/{id}/chat", s.handleChat) mux.HandleFunc("GET /api/agents/{id}/sessions", s.handleListSessions) mux.HandleFunc("GET /api/agents/{id}/sessions/{sid}", s.handleGetSession) + mux.HandleFunc("GET /api/agents/{id}/sessions/{sid}/status", s.handleSessionStatus) // Phase 3: Create & Configure routes mux.HandleFunc("GET /api/wizard/meta", s.handleGetWizardMeta) diff --git a/forge-ui/static/app.js b/forge-ui/static/app.js index 0e7b4dfe..bd971995 100644 --- a/forge-ui/static/app.js +++ b/forge-ui/static/app.js @@ -58,6 +58,16 @@ async function fetchSession(agentId, sessionId) { return res.json(); } +// Reports whether the agent process still considers this task in flight +// ("working"/"submitted") vs settled. Used right after resuming a session +// on page load to detect a turn that was still generating when the +// browser was refreshed away. +async function fetchSessionStatus(agentId, sessionId) { + const res = await fetch(`/api/agents/${agentId}/sessions/${sessionId}/status`); + if (!res.ok) return { state: 'unknown' }; + return res.json(); +} + // ── Phase 3 API Helpers ────────────────────────────────────── async function fetchWizardMeta() { @@ -287,10 +297,15 @@ function useHashRoute() { function parseHash(hash) { const path = hash.replace(/^#\/?/, '') || ''; + // #/agent/{id}/session/{sid} (session segment optional) + const agentSessionMatch = path.match(/^agent\/([^/]+)\/session\/(.+)$/); + if (agentSessionMatch) { + return { page: 'chat', params: { id: agentSessionMatch[1], sessionId: agentSessionMatch[2] } }; + } // #/agent/{id} - const agentMatch = path.match(/^agent\/(.+)$/); + const agentMatch = path.match(/^agent\/([^/]+)$/); if (agentMatch) { - return { page: 'chat', params: { id: agentMatch[1] } }; + return { page: 'chat', params: { id: agentMatch[1], sessionId: null } }; } // #/create if (path === 'create') return { page: 'create', params: {} }; @@ -455,10 +470,10 @@ function formatToolContent(content) { // ── Chat Stream Hook ───────────────────────────────────────── -function useChatStream(agentId) { +function useChatStream(agentId, initialSessionId) { const [messages, setMessages] = useState([]); const [streaming, setStreaming] = useState(false); - const [sessionId, setSessionId] = useState(null); + const [sessionId, setSessionId] = useState(initialSessionId || null); const abortRef = useRef(null); const loadSession = useCallback(async (sid) => { @@ -1027,17 +1042,67 @@ function MessageBubble({ message }) { // ── Chat Page Component ────────────────────────────────────── -function ChatPage({ agentId, agents }) { +function ChatPage({ agentId, initialSessionId, agents }) { const agent = useMemo(() => agents.find(a => a.id === agentId), [agents, agentId]); const isRunning = agent && (agent.status === 'running' || agent.status === 'starting'); - const { messages, streaming, sessionId, sendMessage, loadSession, newSession, cancel } = useChatStream(agentId); + const { messages, streaming, sessionId, sendMessage, loadSession, newSession, cancel } = useChatStream(agentId, initialSessionId); const [sessions, setSessions] = useState([]); const [inputText, setInputText] = useState(''); const messagesEndRef = useRef(null); const messagesContainerRef = useRef(null); const userScrolledUp = useRef(false); const textareaRef = useRef(null); + const loadedInitialRef = useRef(false); + + // Resume the session named in the URL (e.g. after a page reload) instead + // of starting from a blank chat. Runs once per agent mount; a fresh + // initialSessionId only arrives via a full navigation, which remounts + // this component with a new agentId/initialSessionId pair anyway. + // + // After loading, check whether the agent process still considers this + // task in flight — the reload may have interrupted a still-generating + // response. If so, poll briefly and reload the session once it settles, + // so the finished answer appears without the user having to guess + // whether their message went through. + useEffect(() => { + if (!initialSessionId || loadedInitialRef.current) return; + loadedInitialRef.current = true; + let cancelled = false; + + (async () => { + await loadSession(initialSessionId); + + const maxAttempts = 20; // ~30s at the poll interval below + for (let attempt = 0; attempt < maxAttempts && !cancelled; attempt++) { + let status; + try { + status = await fetchSessionStatus(agentId, initialSessionId); + } catch { + break; + } + if (status.state !== 'working' && status.state !== 'submitted') { + if (attempt > 0 && !cancelled) await loadSession(initialSessionId); + break; + } + await new Promise(r => setTimeout(r, 1500)); + } + })(); + + return () => { cancelled = true; }; + }, [agentId, initialSessionId, loadSession]); + + // Keep the URL's session segment in sync with the active session so a + // reload always resumes the same conversation instead of silently + // forking a new one. Uses replaceState (not the hash router's + // navigate()) so this doesn't add reload-noise to browser history. + useEffect(() => { + if (!agentId) return; + const target = sessionId ? `#/agent/${agentId}/session/${sessionId}` : `#/agent/${agentId}`; + if (location.hash !== target) { + history.replaceState(null, '', target); + } + }, [agentId, sessionId]); // Load sessions on mount useEffect(() => { @@ -3296,7 +3361,7 @@ function App() { const renderPage = () => { switch (route.page) { case 'chat': - return html`<${ChatPage} agentId=${route.params.id} agents=${agents} />`; + return html`<${ChatPage} agentId=${route.params.id} initialSessionId=${route.params.sessionId} agents=${agents} />`; case 'create': return html`<${CreatePage} />`; case 'config': From 34b52bd1a3f1a2bebdcc841e7e59d83a8410d243 Mon Sep 17 00:00:00 2001 From: Utkarsh Katiyar Date: Thu, 20 Aug 2026 10:11:41 +0530 Subject: [PATCH 2/2] fix(tests): ensure context cancellation is properly handled in session accumulation test --- forge-core/runtime/loop_refresh_persistence_test.go | 1 + 1 file changed, 1 insertion(+) diff --git a/forge-core/runtime/loop_refresh_persistence_test.go b/forge-core/runtime/loop_refresh_persistence_test.go index 0f1f3a81..8099b796 100644 --- a/forge-core/runtime/loop_refresh_persistence_test.go +++ b/forge-core/runtime/loop_refresh_persistence_test.go @@ -116,6 +116,7 @@ func TestLLMExecutor_RefreshThenRetry_SessionAccumulates(t *testing.T) { }) var ctx context.Context ctx, cancel = context.WithCancel(context.Background()) + defer cancel() _, err = cancelExec.Execute(ctx, &a2a.Task{ID: taskID}, &a2a.Message{Role: a2a.MessageRoleUser, Parts: []a2a.Part{a2a.NewTextPart("first message")}},