diff --git a/cmd/keyoku-server/handlers.go b/cmd/keyoku-server/handlers.go index ab309be..a9acead 100644 --- a/cmd/keyoku-server/handlers.go +++ b/cmd/keyoku-server/handlers.go @@ -30,6 +30,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"` @@ -51,6 +52,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"` diff --git a/cmd/keyoku-server/handlers_memory.go b/cmd/keyoku-server/handlers_memory.go index 40ef2a8..bb311ff 100644 --- a/cmd/keyoku-server/handlers_memory.go +++ b/cmd/keyoku-server/handlers_memory.go @@ -4,6 +4,7 @@ package main import ( + "context" "fmt" "net/http" "sort" @@ -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 @@ -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 @@ -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