claude-code-proxy/proxy/internal/service/conversation_test.go
sid 8e550b9785 Local fork: hardening + ops improvements (timeout knob, demotion, /livez, drain)
This commit captures both the prior accumulated work-in-progress
(framework migration web/→svelte/, postgres storage, conversation
viewer, dashboard auth, OpenAPI spec, integration tests) AND today's
operational improvements layered on top. History wasn't checkpointed
incrementally; happy to split it via interactive rebase if a reviewer
wants smaller commits.

Today's changes (in addition to the older WIP):

1. Configurable upstream response-header timeout
   - ANTHROPIC_RESPONSE_HEADER_TIMEOUT env (default 300s)
   - Replaces hardcoded 300s in provider/anthropic.go that was firing
     on opus + 1M-context + extended thinking non-streaming requests
   - Files: internal/config/config.go, internal/provider/anthropic.go

2. Structured forward-error diagnostic logging
   - When a forward to Anthropic fails, log a single key=value line
     with request_id, model, stream, body_bytes, has_thinking,
     anthropic_beta, query, elapsed, ctx_err — alongside the existing
     human-readable error line for back-compat
   - Files: internal/handler/handlers.go (logForwardFailure)

3. Full SSE protocol passthrough + Flusher fix
   - handler/handlers.go: forward all SSE lines verbatim (event:, id:,
     retry:, : comments, blank-line terminators), not only data:.
     Previous code produced malformed SSE for strict parsers.
   - middleware/logging.go: explicit Flush() method on responseWriter.
     Embedding http.ResponseWriter (interface) does not auto-promote
     Flush(), so every w.(http.Flusher) check in the streaming
     handler was returning ok=false and SSE writes buffered in net/http
     until the body closed.

4. Non-streaming → streaming demotion (feature-flagged)
   - ANTHROPIC_DEMOTE_NONSTREAMING env (default false)
   - When enabled and the routed provider is anthropic, force stream=true
     upstream for clients that asked for stream=false. Receive SSE,
     accumulate via accumulateSSEToMessage (handles text, tool_use with
     partial_json reassembly, thinking, signature, citations_delta,
     usage merge), and synthesize a single non-streaming JSON response.
   - Eliminates the ResponseHeaderTimeout class of failure entirely.
   - Body rewrite uses json.Decoder + UseNumber() to preserve integer
     precision in unknown nested fields (tool inputs from prior turns).
   - Files: internal/config/config.go, internal/handler/handlers.go,
     cmd/proxy/main.go, cmd/proxy/main_test.go

5. Live operational state: /livez gauge + graceful drain
   - New internal/runtime package: atomic in-flight counter + draining flag
   - New middleware/inflight.go: increments runtime gauge, applied to
     /v1/* subrouter so Messages, ChatCompletions, and ProxyPassthrough
     are all counted
   - /v1/* moved to a gorilla/mux subrouter so the InFlight middleware
     applies surgically; /health, /livez, /openapi.* remain on parent
     router (unauthenticated, uncounted)
   - Health handler returns 503 draining when runtime.IsDraining() is
     true, so Traefik stops routing to a slot before drain begins
   - New /livez handler returns {status, in_flight, draining, timestamp}
   - SIGTERM handler in main.go: SetDraining(true), poll for in_flight==0
     with 32-min ceiling and 1s tick (logs every 10s), then srv.Shutdown
   - Auth bypass list extended with /livez
   - Files: internal/runtime/runtime.go (new),
     internal/middleware/inflight.go (new),
     internal/middleware/auth.go,
     internal/handler/handlers.go (Health, Livez, runtime import),
     cmd/proxy/main.go (subrouter, drain loop)

6. OpenAPI spec updates
   - Document Health 503 response and new DrainingResponse schema
   - Add /livez path with LivezResponse schema
   - Files: internal/handler/openapi.go

Verified: go build ./... clean, go test ./... all pass, go vet clean.
Three rounds of codex peer review across changes 1-5; all feedback
addressed (citations_delta, json.Number precision, drain-loop logging
via lastLog timestamp, PathPrefix tightened to "/v1/").
2026-05-02 15:15:58 -06:00

114 lines
3.6 KiB
Go

package service
import (
"os"
"path/filepath"
"testing"
)
func TestConversationServiceAllowsNestedProjectPaths(t *testing.T) {
root := t.TempDir()
projectDir := filepath.Join(root, "team", "app")
if err := os.MkdirAll(projectDir, 0o755); err != nil {
t.Fatalf("MkdirAll() error = %v", err)
}
sessionPath := filepath.Join(projectDir, "session.jsonl")
if err := os.WriteFile(sessionPath, []byte(
`{"timestamp":"2026-03-19T12:00:00Z","type":"user","message":"hello"}`+"\n"+
`{"timestamp":"2026-03-19T12:00:01Z","type":"assistant","message":{"model":"claude-opus-4-6","role":"assistant","content":[{"type":"text","text":"hi"}]}}`+"\n",
), 0o600); err != nil {
t.Fatalf("WriteFile() error = %v", err)
}
svc := &conversationService{claudeProjectsPath: root}
conversation, err := svc.GetConversation("team/app", "session")
if err != nil {
t.Fatalf("GetConversation() error = %v", err)
}
if conversation.SessionID != "session" {
t.Fatalf("expected session ID %q, got %q", "session", conversation.SessionID)
}
if conversation.ProjectPath != "team/app" {
t.Fatalf("expected project path %q, got %q", "team/app", conversation.ProjectPath)
}
if len(conversation.Messages) != 2 {
t.Fatalf("expected 2 messages, got %d", len(conversation.Messages))
}
if conversation.Model != "claude-opus-4-6" {
t.Fatalf("expected model %q, got %q", "claude-opus-4-6", conversation.Model)
}
conversations, err := svc.GetConversationsByProject("team/app")
if err != nil {
t.Fatalf("GetConversationsByProject() error = %v", err)
}
if len(conversations) != 1 {
t.Fatalf("expected 1 conversation, got %d", len(conversations))
}
}
func TestConversationServiceRejectsTraversalPaths(t *testing.T) {
root := t.TempDir()
projectDir := filepath.Join(root, "team", "app")
if err := os.MkdirAll(projectDir, 0o755); err != nil {
t.Fatalf("MkdirAll() error = %v", err)
}
sessionPath := filepath.Join(projectDir, "session.jsonl")
if err := os.WriteFile(sessionPath, []byte(`{"timestamp":"2026-03-19T12:00:00Z","type":"user","message":"hello"}`+"\n"), 0o600); err != nil {
t.Fatalf("WriteFile() error = %v", err)
}
svc := &conversationService{claudeProjectsPath: root}
if _, err := svc.GetConversation("../outside", "session"); err == nil {
t.Fatal("expected traversal project path to be rejected")
}
if _, err := svc.GetConversation("team/app", "../session"); err == nil {
t.Fatal("expected traversal session ID to be rejected")
}
if _, err := svc.GetConversationsByProject("../../outside"); err == nil {
t.Fatal("expected traversal project listing to be rejected")
}
}
func TestConversationServiceRejectsSymlinkEscapes(t *testing.T) {
root := t.TempDir()
projectDir := filepath.Join(root, "team")
if err := os.MkdirAll(projectDir, 0o755); err != nil {
t.Fatalf("MkdirAll() error = %v", err)
}
outsideDir := filepath.Join(t.TempDir(), "outside")
if err := os.MkdirAll(outsideDir, 0o755); err != nil {
t.Fatalf("MkdirAll() error = %v", err)
}
if err := os.WriteFile(filepath.Join(outsideDir, "session.jsonl"), []byte(`{"timestamp":"2026-03-19T12:00:00Z","type":"user","message":"hello"}`+"\n"), 0o600); err != nil {
t.Fatalf("WriteFile() error = %v", err)
}
linkPath := filepath.Join(projectDir, "app")
if err := os.Symlink(outsideDir, linkPath); err != nil {
t.Skipf("symlink not supported in this environment: %v", err)
}
svc := &conversationService{claudeProjectsPath: root}
if _, err := svc.GetConversation("team/app", "session"); err == nil {
t.Fatal("expected symlink escape to be rejected")
}
if _, err := svc.GetConversationsByProject("team/app"); err == nil {
t.Fatal("expected symlink project listing to be rejected")
}
}