Skip to content

feat(api): map timeout_ms exhaustion to 504 request_timeout (builds on #40) #43

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
tyejcoleman wants to merge 2 commits into
base: main
Choose a base branch
from
Open

feat(api): map timeout_ms exhaustion to 504 request_timeout (builds on #40) #43

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7a64d5a
feat(api): support timeout_ms on remember and search
moltar-bot Apr 1, 2026
d38bf26
feat(api): map timeout_ms exhaustion to 504 request_timeout
tyejcoleman Apr 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 48 additions & 3 deletions cmd/keyoku-server/handlers.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,9 @@
package main

import (
"context"
"encoding/json"
"errors"
"log"
"net/http"
"strconv"
Expand All @@ -30,6 +32,7 @@ func NewHandlers(k *keyoku.Keyoku, hub *SSEHub) *Handlers {
type rememberRequest struct {
EntityID string `json:"entity_id"`
Content string `json:"content"`
TimeoutMs int `json:"timeout_ms,omitempty"`
SessionID string `json:"session_id,omitempty"`
AgentID string `json:"agent_id,omitempty"`
Source string `json:"source,omitempty"`
Expand All @@ -51,6 +54,7 @@ type rememberResponse struct {
type searchRequest struct {
EntityID string `json:"entity_id"`
Query string `json:"query"`
TimeoutMs int `json:"timeout_ms,omitempty"`
Limit int `json:"limit,omitempty"`
Mode string `json:"mode,omitempty"`
AgentID string `json:"agent_id,omitempty"`
Expand Down Expand Up @@ -360,11 +364,52 @@ func writeInternalError(w http.ResponseWriter, err error) {
writeInternalErrorWithContext(w, "", err)
}

func writeInternalErrorWithContext(w http.ResponseWriter, context string, err error) {
if context == "" {
// isDeadlineExceeded returns true for wrapped context.DeadlineExceeded and also
// catches providers that return a non-wrapping string (older SDKs). Matches on
// the stable Go-standard message rather than arbitrary provider text.
func isDeadlineExceeded(err error) bool {
if errors.Is(err, context.DeadlineExceeded) {
return true
}
return strings.Contains(err.Error(), "context deadline exceeded")
}

func isCanceled(err error) bool {
if errors.Is(err, context.Canceled) {
return true
}
return strings.Contains(err.Error(), "context canceled")

Copilot AI Apr 24, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

isCanceled uses strings.Contains(err.Error(), "context canceled"). That can misclassify unrelated errors that merely include this substring, and (via the early return in writeInternalErrorWithContext) can lead to an empty/default 200 response. Consider restricting cancellation detection to errors.Is(err, context.Canceled) (or an exact-match fallback) to avoid swallowing real failures.

Suggested change
return strings.Contains(err.Error(), "context canceled")
return err != nil && err.Error() == "context canceled"

Copilot uses AI. Check for mistakes.
}

func writeInternalErrorWithContext(w http.ResponseWriter, ctxLabel string, err error) {
// Caller-controlled budget exhaustion (timeout_ms or upstream deadline) maps
// to 504 with a stable code so clients can distinguish it from server faults.
if isDeadlineExceeded(err) {
if ctxLabel == "" {
log.Printf("INFO: request timed out: %v", err)
} else {
log.Printf("INFO [%s]: request timed out: %v", ctxLabel, err)
}
writeErrorCode(w, http.StatusGatewayTimeout, "request exceeded timeout budget", "request_timeout", true)
return
}

// Client went away (disconnect / upstream cancellation). Nothing useful to
// send back; just log and return without touching the response writer so
// the connection isn't kept alive writing to a dead peer.
if isCanceled(err) {
if ctxLabel == "" {
log.Printf("INFO: request canceled by client: %v", err)
} else {
log.Printf("INFO [%s]: request canceled by client: %v", ctxLabel, err)
}
return
}

if ctxLabel == "" {
log.Printf("ERROR: %v", err)
} else {
log.Printf("ERROR [%s]: %v", context, err)
log.Printf("ERROR [%s]: %v", ctxLabel, err)
}

errStr := err.Error()
Expand Down
58 changes: 58 additions & 0 deletions cmd/keyoku-server/handlers_errors_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,10 @@
package main

import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"net/http/httptest"
"testing"
Expand Down Expand Up @@ -55,6 +57,62 @@ func TestWriteInternalErrorWithContext_GenericInternal(t *testing.T) {
}
}

func TestWriteInternalErrorWithContext_DeadlineExceededMapsTo504(t *testing.T) {
rr := httptest.NewRecorder()

// Simulate the wrapping pattern used by engine_add.go and llm providers.
wrapped := fmt.Errorf("extraction failed: %w", context.DeadlineExceeded)
writeInternalErrorWithContext(rr, "remember", wrapped)

if rr.Code != http.StatusGatewayTimeout {
t.Fatalf("status = %d, want %d", rr.Code, http.StatusGatewayTimeout)
}

var body apiErrorResponse
if err := json.Unmarshal(rr.Body.Bytes(), &body); err != nil {
t.Fatalf("unmarshal response: %v", err)
}
if body.Code != "request_timeout" {
t.Fatalf("code = %q, want request_timeout", body.Code)
}
if !body.Retryable {
t.Fatalf("retryable = false, want true")
}
}

func TestWriteInternalErrorWithContext_DeadlineExceededStringFallback(t *testing.T) {
rr := httptest.NewRecorder()

// Older provider SDKs may return the stdlib text without wrapping the sentinel.
writeInternalErrorWithContext(rr, "search", errors.New("provider call failed: context deadline exceeded"))

if rr.Code != http.StatusGatewayTimeout {
t.Fatalf("status = %d, want %d", rr.Code, http.StatusGatewayTimeout)
}
var body apiErrorResponse
if err := json.Unmarshal(rr.Body.Bytes(), &body); err != nil {
t.Fatalf("unmarshal response: %v", err)
}
if body.Code != "request_timeout" {
t.Fatalf("code = %q, want request_timeout", body.Code)
}
}

func TestWriteInternalErrorWithContext_CanceledWritesNoResponse(t *testing.T) {
rr := httptest.NewRecorder()

wrapped := fmt.Errorf("extraction failed: %w", context.Canceled)
writeInternalErrorWithContext(rr, "remember", wrapped)

// Default recorder status is 200; no WriteHeader should have been called.
if rr.Code != http.StatusOK {
t.Fatalf("status = %d, want default 200 (no header written), got otherwise", rr.Code)
}
if rr.Body.Len() != 0 {
t.Fatalf("body = %q, want empty — client is gone, we should not write", rr.Body.String())
}
}

func TestWriteInternalErrorWithContext_SimilarityPrefixWithoutHNSWIs500(t *testing.T) {
rr := httptest.NewRecorder()

Expand Down
32 changes: 30 additions & 2 deletions cmd/keyoku-server/handlers_memory.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
package main

import (
"context"
"fmt"
"net/http"
"sort"
Expand All @@ -14,6 +15,19 @@ import (
"github.com/keyoku-ai/keyoku-engine/storage"
)

func requestTimeout(timeoutMs int) (time.Duration, bool) {
if timeoutMs <= 0 {
return 0, false
}
if timeoutMs < 1000 {
timeoutMs = 1000
}
if timeoutMs > 300000 {
timeoutMs = 300000
}
return time.Duration(timeoutMs) * time.Millisecond, true
}

// HandleRemember extracts and stores memories from content.
func (h *Handlers) HandleRemember(w http.ResponseWriter, r *http.Request) {
var req rememberRequest
Expand Down Expand Up @@ -66,7 +80,14 @@ func (h *Handlers) HandleRemember(w http.ResponseWriter, r *http.Request) {
opts = append(opts, keyoku.WithCreatedAt(t))
}

result, err := h.k.Remember(r.Context(), req.EntityID, req.Content, opts...)
ctx := r.Context()
if timeout, ok := requestTimeout(req.TimeoutMs); ok {
var cancel context.CancelFunc
ctx, cancel = context.WithTimeout(ctx, timeout)
defer cancel()
}

result, err := h.k.Remember(ctx, req.EntityID, req.Content, opts...)
if err != nil {
writeInternalErrorWithContext(w, "remember", err)
return
Expand Down Expand Up @@ -115,7 +136,14 @@ func (h *Handlers) HandleSearch(w http.ResponseWriter, r *http.Request) {
opts = append(opts, keyoku.WithSearchAgentID(req.AgentID))
}

results, err := h.k.Search(r.Context(), req.EntityID, req.Query, opts...)
ctx := r.Context()
if timeout, ok := requestTimeout(req.TimeoutMs); ok {
var cancel context.CancelFunc
ctx, cancel = context.WithTimeout(ctx, timeout)
defer cancel()
}

results, err := h.k.Search(ctx, req.EntityID, req.Query, opts...)
if err != nil {
writeInternalErrorWithContext(w, "search", err)
return
Expand Down
Loading