server/debug: add detached context for debug server (#751)

Author: WineChord

server/debug/server.go

@@ -569,8 +569,8 @@ func (s *Server) handleRun(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, err.Error(), http.StatusBadRequest)
 		return
 	}
-
-	out, err := rn.Run(r.Context(), req.UserID, req.SessionID,
+	ctx := newDetachedContext(r.Context())
+	out, err := rn.Run(ctx, req.UserID, req.SessionID,
 		convertContentToMessage(req.NewMessage))
 	if err != nil {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
@@ -616,7 +616,8 @@ func (s *Server) handleRunSSE(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, err.Error(), http.StatusBadRequest)
 		return
 	}
-	out, err := rn.Run(context.Background(), req.UserID, req.SessionID,
+	ctx := newDetachedContext(r.Context())
+	out, err := rn.Run(ctx, req.UserID, req.SessionID,
 		convertContentToMessage(req.NewMessage))
 	if err != nil {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
@@ -1619,3 +1620,27 @@ func isGraphToolEvent(e *event.Event) bool {
 	_, exists := e.StateDelta[graph.MetadataKeyTool]
 	return exists
 }
+
+// detachedContext wraps a parent context but disables cancellation and
+// deadlines while preserving all values. This allows us to keep trace and
+// logging metadata from the incoming request context without being affected
+// by HTTP‑level timeouts or client disconnects.
+type detachedContext struct {
+	context.Context
+}
+
+func (detachedContext) Deadline() (time.Time, bool) {
+	return time.Time{}, false
+}
+
+func (detachedContext) Done() <-chan struct{} {
+	return nil
+}
+
+func (detachedContext) Err() error {
+	return nil
+}
+
+func newDetachedContext(ctx context.Context) context.Context {
+	return detachedContext{Context: ctx}
+}

server/debug/server_test.go

@@ -89,6 +89,10 @@ type fakeRunner struct {
 	err    error
 }
 
+type ctxCapturingRunner struct {
+	ctx context.Context
+}
+
 func (f *fakeRunner) Run(ctx context.Context, userID string, sessionID string, message model.Message, runOpts ...agent.RunOption) (<-chan *event.Event, error) {
 	if f.err != nil {
 		return nil, f.err
@@ -105,6 +109,23 @@ func (f *fakeRunner) Run(ctx context.Context, userID string, sessionID string, m
 
 func (f *fakeRunner) Close() error { return nil }
 
+func (f *ctxCapturingRunner) Run(
+	ctx context.Context,
+	userID string,
+	sessionID string,
+	message model.Message,
+	runOpts ...agent.RunOption,
+) (<-chan *event.Event, error) {
+	f.ctx = ctx
+	ch := make(chan *event.Event)
+	close(ch)
+	return ch, nil
+}
+
+func (f *ctxCapturingRunner) Close() error {
+	return nil
+}
+
 type flushRecorder struct {
 	*httptest.ResponseRecorder
 }
@@ -1435,6 +1456,82 @@ func TestHandleRunSSE_NonStreaming(t *testing.T) {
 	assert.Contains(t, w.Body.String(), "data: ")
 }
 
+func TestNewDetachedContextPreservesValues(t *testing.T) {
+	type ctxKey struct{}
+	key := ctxKey{}
+	parent, cancel := context.WithCancel(
+		context.WithValue(context.Background(), key, "trace-id"),
+	)
+	cancel()
+
+	ctx := newDetachedContext(parent)
+
+	assert.Equal(t, "trace-id", ctx.Value(key))
+	_, ok := ctx.Deadline()
+	assert.False(t, ok)
+	assert.Nil(t, ctx.Done())
+	assert.Nil(t, ctx.Err())
+
+	err := agent.CheckContextCancelled(ctx)
+	assert.NoError(t, err)
+}
+
+func TestHandleRunSSE_UsesDetachedContext(t *testing.T) {
+	type ctxKey struct{}
+	key := ctxKey{}
+
+	ctxRunner := &ctxCapturingRunner{}
+	server := &Server{
+		agents:     map[string]agent.Agent{},
+		router:     mux.NewRouter(),
+		runners:    map[string]runner.Runner{"app": ctxRunner},
+		sessionSvc: sessioninmemory.NewSessionService(),
+		traces:     map[string]attribute.Set{},
+
+		memoryExporter: newInMemoryExporter(),
+	}
+
+	reqBody := schema.AgentRunRequest{
+		AppName:   "app",
+		UserID:    "user",
+		SessionID: "sess",
+		NewMessage: schema.Content{
+			Role:  "user",
+			Parts: []schema.Part{{Text: "hi"}},
+		},
+		Streaming: true,
+	}
+	body, err := json.Marshal(reqBody)
+	require.NoError(t, err)
+
+	baseCtx, cancel := context.WithCancel(
+		context.WithValue(context.Background(), key, "trace-id"),
+	)
+	cancel()
+
+	req := httptest.NewRequest(
+		http.MethodPost,
+		"/run_sse",
+		bytes.NewReader(body),
+	).WithContext(baseCtx)
+	req.Header.Set("Content-Type", "application/json")
+	w := &flushRecorder{ResponseRecorder: httptest.NewRecorder()}
+
+	server.handleRunSSE(w, req)
+
+	assert.Equal(t, http.StatusOK, w.Code)
+	if assert.NotNil(t, ctxRunner.ctx) {
+		assert.Equal(t, "trace-id", ctxRunner.ctx.Value(key))
+		_, ok := ctxRunner.ctx.Deadline()
+		assert.False(t, ok)
+		assert.Nil(t, ctxRunner.ctx.Done())
+		assert.Nil(t, ctxRunner.ctx.Err())
+
+		err = agent.CheckContextCancelled(ctxRunner.ctx)
+		assert.NoError(t, err)
+	}
+}
+
 func TestServerGetRunnerCache(t *testing.T) {
 	server := New(map[string]agent.Agent{
 		"app": &mockAgent{name: "app"},