feat(server): provider switching API, WS/REST cancel, summarize endpoints
This commit is contained in:
parent
f1b2093b22
commit
49fdc5a189
7 changed files with 463 additions and 28 deletions
|
|
@ -17,6 +17,7 @@ import (
|
||||||
"github.com/crypto/byteclaw/internal/agent"
|
"github.com/crypto/byteclaw/internal/agent"
|
||||||
"github.com/crypto/byteclaw/internal/config"
|
"github.com/crypto/byteclaw/internal/config"
|
||||||
"github.com/crypto/byteclaw/internal/llm"
|
"github.com/crypto/byteclaw/internal/llm"
|
||||||
|
"github.com/crypto/byteclaw/internal/providers"
|
||||||
"github.com/crypto/byteclaw/internal/server"
|
"github.com/crypto/byteclaw/internal/server"
|
||||||
"github.com/crypto/byteclaw/internal/store"
|
"github.com/crypto/byteclaw/internal/store"
|
||||||
"github.com/crypto/byteclaw/internal/tools"
|
"github.com/crypto/byteclaw/internal/tools"
|
||||||
|
|
@ -71,10 +72,16 @@ func run(cfgPath string, log *slog.Logger) error {
|
||||||
reg.Register(sh)
|
reg.Register(sh)
|
||||||
log.Warn("shell tool enabled", "allowed", sc.AllowedCommands)
|
log.Warn("shell tool enabled", "allowed", sc.AllowedCommands)
|
||||||
}
|
}
|
||||||
prov := llm.NewOpenAI(cfg.LLM.BaseURL, cfg.LLM.APIKey, cfg.LLM.Model, cfg.LLM.Temperature, time.Duration(cfg.LLM.TimeoutSec)*time.Second)
|
pm := providers.New(cfg.LLM, os.Getenv, st)
|
||||||
ag := &agent.Agent{LLM: prov, Tools: reg, Store: st, SystemPrompt: cfg.Agent.SystemPrompt,
|
sm := cfg.Agent.Summarize
|
||||||
MaxIterations: cfg.Agent.MaxIterations, HistoryLimit: cfg.Agent.HistoryLimit}
|
ag := &agent.Agent{Tools: reg, Store: st, SystemPrompt: cfg.Agent.SystemPrompt,
|
||||||
srv := &server.Server{Agent: ag, Store: st, Token: cfg.AuthToken, AllowedOrigins: cfg.AllowedOrigins, Log: log,
|
MaxIterations: cfg.Agent.MaxIterations, HistoryLimit: cfg.Agent.HistoryLimit,
|
||||||
|
Summarize: agent.SummarizeConfig{Enabled: sm.Enabled, MaxMessages: sm.MaxMessages, MaxTokens: sm.MaxTokens, KeepRecent: sm.KeepRecent},
|
||||||
|
Resolve: func(ctx context.Context, sid string) (llm.Provider, error) {
|
||||||
|
p, _, err := pm.ForSession(ctx, sid)
|
||||||
|
return p, err
|
||||||
|
}}
|
||||||
|
srv := &server.Server{Agent: ag, Store: st, Token: cfg.AuthToken, AllowedOrigins: cfg.AllowedOrigins, Log: log, Providers: pm,
|
||||||
Info: server.Info{Version: version, Provider: cfg.LLM.Provider, Model: cfg.LLM.Model, BaseURL: cfg.LLM.BaseURL},
|
Info: server.Info{Version: version, Provider: cfg.LLM.Provider, Model: cfg.LLM.Model, BaseURL: cfg.LLM.BaseURL},
|
||||||
RunTimeout: time.Duration(cfg.LLM.TimeoutSec*cfg.Agent.MaxIterations) * time.Second}
|
RunTimeout: time.Duration(cfg.LLM.TimeoutSec*cfg.Agent.MaxIterations) * time.Second}
|
||||||
|
|
||||||
|
|
@ -83,7 +90,7 @@ func run(cfgPath string, log *slog.Logger) error {
|
||||||
defer stop()
|
defer stop()
|
||||||
errc := make(chan error, 1)
|
errc := make(chan error, 1)
|
||||||
go func() {
|
go func() {
|
||||||
log.Info("byteclaw listening", "addr", cfg.Listen, "provider", cfg.LLM.Provider, "model", cfg.LLM.Model, "tools", len(reg.Specs()), "version", version)
|
log.Info("byteclaw listening", "addr", cfg.Listen, "provider", pm.Active().Provider, "model", pm.Active().Model, "tools", len(reg.Specs()), "version", version)
|
||||||
errc <- hs.ListenAndServe()
|
errc <- hs.ListenAndServe()
|
||||||
}()
|
}()
|
||||||
select {
|
select {
|
||||||
|
|
|
||||||
|
|
@ -17,11 +17,23 @@ llm:
|
||||||
model: "" # e.g. MiniMaxAI/MiniMax-M2.7, zai-org/GLM-5.3-Flash
|
model: "" # e.g. MiniMaxAI/MiniMax-M2.7, zai-org/GLM-5.3-Flash
|
||||||
timeout_sec: 120
|
timeout_sec: 120
|
||||||
temperature: 0.7
|
temperature: 0.7
|
||||||
|
# Further providers the app may switch to at runtime (PUT /api/providers/active or per session).
|
||||||
|
# Keys stay on the server and are never sent to the app. Presets also pick up their env key.
|
||||||
|
providers:
|
||||||
|
# dahl: { api_key: "" } # or env DAHL_API_KEY
|
||||||
|
# gonkarouter: { api_key: "" } # or env GONKAROUTER_API_KEY
|
||||||
|
# mylocal: { base_url: "http://127.0.0.1:8000/v1", model: "qwen3" }
|
||||||
|
|
||||||
agent:
|
agent:
|
||||||
system_prompt: "You are ByteClaw, a concise and helpful self-hosted assistant. Use tools when they help."
|
system_prompt: "You are ByteClaw, a concise and helpful self-hosted assistant. Use tools when they help."
|
||||||
max_iterations: 8
|
max_iterations: 8
|
||||||
history_limit: 50
|
history_limit: 50
|
||||||
|
# Summarise long sessions with the LLM (originals stay in SQLite). Manual: POST /api/sessions/{id}/summarize
|
||||||
|
summarize:
|
||||||
|
enabled: true
|
||||||
|
max_messages: 40 # unsummarised messages before summarising
|
||||||
|
max_tokens: 12000 # or estimated tokens (chars/4)
|
||||||
|
keep_recent: 10 # newest messages always kept verbatim
|
||||||
|
|
||||||
tools:
|
tools:
|
||||||
web_fetch:
|
web_fetch:
|
||||||
|
|
|
||||||
|
|
@ -89,7 +89,7 @@ func (a *Agent) summarizeLocked(ctx context.Context, prov llm.Provider, sessionI
|
||||||
resp, err := prov.Chat(ctx, llm.Request{Messages: []llm.Message{
|
resp, err := prov.Chat(ctx, llm.Request{Messages: []llm.Message{
|
||||||
{Role: llm.RoleSystem, Content: summarizePrompt},
|
{Role: llm.RoleSystem, Content: summarizePrompt},
|
||||||
{Role: llm.RoleUser, Content: b.String()},
|
{Role: llm.RoleUser, Content: b.String()},
|
||||||
}}, nil)
|
}}, func(string) {})
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("summarize: %w", err)
|
return nil, fmt.Errorf("summarize: %w", err)
|
||||||
}
|
}
|
||||||
|
|
|
||||||
39
internal/llm/cancel_test.go
Normal file
39
internal/llm/cancel_test.go
Normal file
|
|
@ -0,0 +1,39 @@
|
||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Cancelling the context must abort the in-flight provider HTTP request.
|
||||||
|
func TestOpenAICancelAbortsRequest(t *testing.T) {
|
||||||
|
aborted := make(chan struct{})
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"Hal\"}}]}\n\n")
|
||||||
|
w.(http.Flusher).Flush()
|
||||||
|
<-r.Context().Done()
|
||||||
|
close(aborted)
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
p := NewOpenAI(srv.URL, "k", "m", 0, 30*time.Second)
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
got := ""
|
||||||
|
_, err := p.Chat(ctx, Request{Messages: []Message{{Role: RoleUser, Content: "x"}}}, func(s string) {
|
||||||
|
got += s
|
||||||
|
cancel()
|
||||||
|
})
|
||||||
|
if err == nil || !errors.Is(ctx.Err(), context.Canceled) || got != "Hal" {
|
||||||
|
t.Fatalf("err=%v got=%q", err, got)
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case <-aborted:
|
||||||
|
case <-time.After(3 * time.Second):
|
||||||
|
t.Fatal("provider request was not aborted")
|
||||||
|
}
|
||||||
|
}
|
||||||
146
internal/server/extra.go
Normal file
146
internal/server/extra.go
Normal file
|
|
@ -0,0 +1,146 @@
|
||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"errors"
|
||||||
|
"net/http"
|
||||||
|
|
||||||
|
"github.com/crypto/byteclaw/internal/providers"
|
||||||
|
"github.com/crypto/byteclaw/internal/store"
|
||||||
|
)
|
||||||
|
|
||||||
|
func (s *Server) info(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
in := s.Info
|
||||||
|
if s.Providers != nil {
|
||||||
|
a := s.Providers.Active()
|
||||||
|
in.Provider, in.Model, in.BaseURL = a.Provider, a.Model, s.Providers.BaseURL(a.Provider)
|
||||||
|
}
|
||||||
|
writeJSON(w, 200, in)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Server) needProviders(w http.ResponseWriter) bool {
|
||||||
|
if s.Providers == nil {
|
||||||
|
writeErr(w, http.StatusNotImplemented, "provider switching not enabled")
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
func providerErr(w http.ResponseWriter, err error) {
|
||||||
|
switch {
|
||||||
|
case errors.Is(err, providers.ErrUnknown):
|
||||||
|
writeErr(w, http.StatusNotFound, err.Error())
|
||||||
|
default:
|
||||||
|
writeErr(w, http.StatusBadRequest, err.Error())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type providerReq struct {
|
||||||
|
Provider string `json:"provider"`
|
||||||
|
Model string `json:"model"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// GET /api/providers -> {"active":{...},"providers":[...]} (never contains keys).
|
||||||
|
func (s *Server) listProviders(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
if !s.needProviders(w) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
writeJSON(w, 200, map[string]any{"active": s.Providers.Active(), "providers": s.Providers.List()})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Server) setActiveProvider(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if !s.needProviders(w) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
var b providerReq
|
||||||
|
if !decode(w, r, &b) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
a, err := s.Providers.SetActive(r.Context(), b.Provider, b.Model)
|
||||||
|
if err != nil {
|
||||||
|
providerErr(w, err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
writeJSON(w, 200, a)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Server) listModels(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if !s.needProviders(w) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
ids, err := s.Providers.Models(r.Context(), r.PathValue("name"))
|
||||||
|
if errors.Is(err, providers.ErrUnknown) {
|
||||||
|
providerErr(w, err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
s.Log.Warn("list models", "err", err)
|
||||||
|
writeErr(w, http.StatusBadGateway, "provider does not offer a model list")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
writeJSON(w, 200, map[string]any{"models": ids})
|
||||||
|
}
|
||||||
|
|
||||||
|
// PUT /api/sessions/{id}/provider: per-session override; empty provider clears it.
|
||||||
|
func (s *Server) setSessionProvider(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if !s.needProviders(w) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
var b providerReq
|
||||||
|
if !decode(w, r, &b) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
id := r.PathValue("id")
|
||||||
|
if _, err := s.Store.GetSession(r.Context(), id); err != nil {
|
||||||
|
s.storeErr(w, err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
var a providers.Active
|
||||||
|
if b.Provider != "" {
|
||||||
|
var err error
|
||||||
|
if a, err = s.Providers.Validate(b.Provider, b.Model); err != nil {
|
||||||
|
providerErr(w, err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if err := s.Store.SetSessionProvider(r.Context(), id, a.Provider, a.Model); err != nil {
|
||||||
|
s.storeErr(w, err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
writeJSON(w, 200, a)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Server) cancel(w http.ResponseWriter, r *http.Request) {
|
||||||
|
writeJSON(w, 200, map[string]bool{"cancelled": s.Agent.Cancel(r.PathValue("id"))})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Server) summarize(w http.ResponseWriter, r *http.Request) {
|
||||||
|
id := r.PathValue("id")
|
||||||
|
if _, err := s.Store.GetSession(r.Context(), id); err != nil {
|
||||||
|
s.storeErr(w, err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
sm, err := s.Agent.SummarizeNow(r.Context(), id)
|
||||||
|
if err != nil {
|
||||||
|
s.Log.Warn("summarize", "err", err)
|
||||||
|
writeErr(w, http.StatusBadGateway, err.Error())
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if sm == nil {
|
||||||
|
writeJSON(w, 200, map[string]any{"summary": nil, "note": "nothing to summarize yet"})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
writeJSON(w, 200, map[string]any{"summary": sm})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Server) getSummary(w http.ResponseWriter, r *http.Request) {
|
||||||
|
sm, err := s.Store.LatestSummary(r.Context(), r.PathValue("id"))
|
||||||
|
if err != nil {
|
||||||
|
if errors.Is(err, store.ErrNotFound) {
|
||||||
|
writeErr(w, http.StatusNotFound, "no summary")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
s.storeErr(w, err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
writeJSON(w, 200, sm)
|
||||||
|
}
|
||||||
180
internal/server/extra_test.go
Normal file
180
internal/server/extra_test.go
Normal file
|
|
@ -0,0 +1,180 @@
|
||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"io"
|
||||||
|
"log/slog"
|
||||||
|
"net/http/httptest"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/coder/websocket"
|
||||||
|
"github.com/coder/websocket/wsjson"
|
||||||
|
|
||||||
|
"github.com/crypto/byteclaw/internal/agent"
|
||||||
|
"github.com/crypto/byteclaw/internal/config"
|
||||||
|
"github.com/crypto/byteclaw/internal/llm"
|
||||||
|
"github.com/crypto/byteclaw/internal/providers"
|
||||||
|
"github.com/crypto/byteclaw/internal/store"
|
||||||
|
"github.com/crypto/byteclaw/internal/tools"
|
||||||
|
)
|
||||||
|
|
||||||
|
// slow streams a first chunk then blocks until the context is cancelled.
|
||||||
|
type slow struct{}
|
||||||
|
|
||||||
|
func (slow) Chat(ctx context.Context, _ llm.Request, d llm.DeltaFunc) (llm.Response, error) {
|
||||||
|
d("halb")
|
||||||
|
<-ctx.Done()
|
||||||
|
return llm.Response{}, ctx.Err()
|
||||||
|
}
|
||||||
|
|
||||||
|
func setupFull(t *testing.T, p llm.Provider) (*httptest.Server, *store.Store) {
|
||||||
|
st, err := store.Open(filepath.Join(t.TempDir(), "s.db"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
t.Cleanup(func() { st.Close() })
|
||||||
|
cfg := config.LLM{Provider: "openai", BaseURL: "https://api.openai.com/v1", Model: "gpt-4o-mini", APIKey: "sk-SECRET-openai"}
|
||||||
|
pm := providers.New(cfg, func(k string) string {
|
||||||
|
if k == "DAHL_API_KEY" {
|
||||||
|
return "dahl-SECRET"
|
||||||
|
}
|
||||||
|
return ""
|
||||||
|
}, st)
|
||||||
|
pm.New = func(string, string, string) llm.Provider { return p }
|
||||||
|
ag := &agent.Agent{Store: st, Tools: tools.NewRegistry(), MaxIterations: 3, Summarize: agent.SummarizeConfig{KeepRecent: 1}}
|
||||||
|
ag.Resolve = func(ctx context.Context, sid string) (llm.Provider, error) {
|
||||||
|
pr, _, err := pm.ForSession(ctx, sid)
|
||||||
|
return pr, err
|
||||||
|
}
|
||||||
|
s := &Server{Store: st, Token: tok, Log: slog.New(slog.NewTextHandler(io.Discard, nil)), RunTimeout: 5 * time.Second, Agent: ag, Providers: pm}
|
||||||
|
ts := httptest.NewServer(s.Handler())
|
||||||
|
t.Cleanup(ts.Close)
|
||||||
|
return ts, st
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProviderEndpoints(t *testing.T) {
|
||||||
|
ts, st := setupFull(t, echo{})
|
||||||
|
if r := do(t, "GET", ts.URL+"/api/providers", "", nil); r.StatusCode != 401 {
|
||||||
|
t.Fatal("providers must need token")
|
||||||
|
}
|
||||||
|
r := do(t, "GET", ts.URL+"/api/providers", "", authH)
|
||||||
|
b, _ := io.ReadAll(r.Body)
|
||||||
|
if r.StatusCode != 200 || strings.Contains(string(b), "SECRET") || strings.Contains(string(b), "api_key") {
|
||||||
|
t.Fatalf("%d leaked/failed: %s", r.StatusCode, b)
|
||||||
|
}
|
||||||
|
if !strings.Contains(string(b), `"name":"dahl"`) {
|
||||||
|
t.Fatalf("%s", b)
|
||||||
|
}
|
||||||
|
if r := do(t, "PUT", ts.URL+"/api/providers/active", `{"provider":"openrouter"}`, authH); r.StatusCode != 400 {
|
||||||
|
t.Fatal("unconfigured provider must be rejected", r.StatusCode)
|
||||||
|
}
|
||||||
|
if r := do(t, "PUT", ts.URL+"/api/providers/active", `{"provider":"nope"}`, authH); r.StatusCode != 404 {
|
||||||
|
t.Fatal(r.StatusCode)
|
||||||
|
}
|
||||||
|
if r := do(t, "PUT", ts.URL+"/api/providers/active", `{"provider":"dahl","model":"zai-org/GLM-5.3-Flash"}`, authH); r.StatusCode != 200 {
|
||||||
|
t.Fatal(r.StatusCode)
|
||||||
|
}
|
||||||
|
var in Info
|
||||||
|
json.NewDecoder(do(t, "GET", ts.URL+"/api/info", "", authH).Body).Decode(&in)
|
||||||
|
if in.Provider != "dahl" || in.Model != "zai-org/GLM-5.3-Flash" || in.BaseURL != "https://inference.dahl.global/v1" {
|
||||||
|
t.Fatalf("%+v", in)
|
||||||
|
}
|
||||||
|
ss, _ := st.CreateSession(context.Background(), "x")
|
||||||
|
if r := do(t, "PUT", ts.URL+"/api/sessions/"+ss.ID+"/provider", `{"provider":"openai"}`, authH); r.StatusCode != 200 {
|
||||||
|
t.Fatal(r.StatusCode)
|
||||||
|
}
|
||||||
|
got, _ := st.GetSession(context.Background(), ss.ID)
|
||||||
|
if got.Provider != "openai" || got.Model != "gpt-4o-mini" {
|
||||||
|
t.Fatalf("%+v", got)
|
||||||
|
}
|
||||||
|
if r := do(t, "PUT", ts.URL+"/api/sessions/missing/provider", `{"provider":"openai"}`, authH); r.StatusCode != 404 {
|
||||||
|
t.Fatal(r.StatusCode)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWebSocketCancel(t *testing.T) {
|
||||||
|
ts, st := setupFull(t, slow{})
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
c, _, err := websocket.Dial(ctx, "ws"+strings.TrimPrefix(ts.URL, "http")+"/api/ws", &websocket.DialOptions{HTTPHeader: map[string][]string{"Authorization": {"Bearer " + tok}}})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer c.CloseNow()
|
||||||
|
wsjson.Write(ctx, c, chatReq{Type: "message", Text: "lange antwort"})
|
||||||
|
var sid string
|
||||||
|
for {
|
||||||
|
var e agent.Event
|
||||||
|
if err := wsjson.Read(ctx, c, &e); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
switch e.Type {
|
||||||
|
case "session":
|
||||||
|
sid = e.Session
|
||||||
|
case "delta":
|
||||||
|
wsjson.Write(ctx, c, chatReq{Type: "cancel", SessionID: sid})
|
||||||
|
case "cancelled":
|
||||||
|
if e.Text != "halb" {
|
||||||
|
t.Fatalf("%+v", e)
|
||||||
|
}
|
||||||
|
msgs, _ := st.Messages(ctx, sid, 0)
|
||||||
|
if len(msgs) != 2 || !msgs[1].Cancelled || msgs[1].Content != "halb" {
|
||||||
|
t.Fatalf("%+v", msgs)
|
||||||
|
}
|
||||||
|
return
|
||||||
|
case "error", "done":
|
||||||
|
t.Fatalf("unexpected %+v", e)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRESTCancelAndSummarize(t *testing.T) {
|
||||||
|
ts, st := setupFull(t, slow{})
|
||||||
|
ss, _ := st.CreateSession(context.Background(), "x")
|
||||||
|
done := make(chan string, 1)
|
||||||
|
go func() {
|
||||||
|
r := do(t, "POST", ts.URL+"/api/chat", `{"session_id":"`+ss.ID+`","text":"hi"}`, authH)
|
||||||
|
b, _ := io.ReadAll(r.Body)
|
||||||
|
done <- string(b)
|
||||||
|
}()
|
||||||
|
deadline := time.Now().Add(3 * time.Second)
|
||||||
|
for {
|
||||||
|
r := do(t, "POST", ts.URL+"/api/sessions/"+ss.ID+"/cancel", "", authH)
|
||||||
|
var b map[string]bool
|
||||||
|
json.NewDecoder(r.Body).Decode(&b)
|
||||||
|
if b["cancelled"] {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
if time.Now().After(deadline) {
|
||||||
|
t.Fatal("never running")
|
||||||
|
}
|
||||||
|
time.Sleep(10 * time.Millisecond)
|
||||||
|
}
|
||||||
|
if body := <-done; !strings.Contains(body, `"cancelled":true`) {
|
||||||
|
t.Fatal(body)
|
||||||
|
}
|
||||||
|
if r := do(t, "GET", ts.URL+"/api/sessions/"+ss.ID+"/summary", "", authH); r.StatusCode != 404 {
|
||||||
|
t.Fatal(r.StatusCode)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestManualSummarize(t *testing.T) {
|
||||||
|
ts, st := setupFull(t, echo{})
|
||||||
|
ctx := context.Background()
|
||||||
|
ss, _ := st.CreateSession(ctx, "x")
|
||||||
|
for _, q := range []string{"a", "b"} {
|
||||||
|
do(t, "POST", ts.URL+"/api/chat", `{"session_id":"`+ss.ID+`","text":"`+q+`"}`, authH)
|
||||||
|
}
|
||||||
|
r := do(t, "POST", ts.URL+"/api/sessions/"+ss.ID+"/summarize", "", authH)
|
||||||
|
b, _ := io.ReadAll(r.Body)
|
||||||
|
if r.StatusCode != 200 || !strings.Contains(string(b), `"upto_message_id"`) {
|
||||||
|
t.Fatalf("%d %s", r.StatusCode, b)
|
||||||
|
}
|
||||||
|
if r := do(t, "GET", ts.URL+"/api/sessions/"+ss.ID+"/summary", "", authH); r.StatusCode != 200 {
|
||||||
|
t.Fatal(r.StatusCode)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -10,6 +10,7 @@ import (
|
||||||
"net/http"
|
"net/http"
|
||||||
"net/url"
|
"net/url"
|
||||||
"strings"
|
"strings"
|
||||||
|
"sync"
|
||||||
"time"
|
"time"
|
||||||
"unicode/utf8"
|
"unicode/utf8"
|
||||||
|
|
||||||
|
|
@ -17,6 +18,7 @@ import (
|
||||||
"github.com/coder/websocket/wsjson"
|
"github.com/coder/websocket/wsjson"
|
||||||
|
|
||||||
"github.com/crypto/byteclaw/internal/agent"
|
"github.com/crypto/byteclaw/internal/agent"
|
||||||
|
"github.com/crypto/byteclaw/internal/providers"
|
||||||
"github.com/crypto/byteclaw/internal/store"
|
"github.com/crypto/byteclaw/internal/store"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -31,6 +33,8 @@ type Server struct {
|
||||||
RunTimeout time.Duration
|
RunTimeout time.Duration
|
||||||
// Info is returned by GET /api/info (never contains secrets).
|
// Info is returned by GET /api/info (never contains secrets).
|
||||||
Info Info
|
Info Info
|
||||||
|
// Providers enables runtime provider switching (optional).
|
||||||
|
Providers *providers.Manager
|
||||||
}
|
}
|
||||||
|
|
||||||
type Info struct {
|
type Info struct {
|
||||||
|
|
@ -44,7 +48,14 @@ func (s *Server) Handler() http.Handler {
|
||||||
mux := http.NewServeMux()
|
mux := http.NewServeMux()
|
||||||
mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, _ *http.Request) { writeJSON(w, 200, map[string]string{"status": "ok"}) })
|
mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, _ *http.Request) { writeJSON(w, 200, map[string]string{"status": "ok"}) })
|
||||||
api := http.NewServeMux()
|
api := http.NewServeMux()
|
||||||
api.HandleFunc("GET /api/info", func(w http.ResponseWriter, _ *http.Request) { writeJSON(w, 200, s.Info) })
|
api.HandleFunc("GET /api/info", s.info)
|
||||||
|
api.HandleFunc("GET /api/providers", s.listProviders)
|
||||||
|
api.HandleFunc("PUT /api/providers/active", s.setActiveProvider)
|
||||||
|
api.HandleFunc("GET /api/providers/{name}/models", s.listModels)
|
||||||
|
api.HandleFunc("PUT /api/sessions/{id}/provider", s.setSessionProvider)
|
||||||
|
api.HandleFunc("POST /api/sessions/{id}/cancel", s.cancel)
|
||||||
|
api.HandleFunc("POST /api/sessions/{id}/summarize", s.summarize)
|
||||||
|
api.HandleFunc("GET /api/sessions/{id}/summary", s.getSummary)
|
||||||
api.HandleFunc("GET /api/sessions", s.listSessions)
|
api.HandleFunc("GET /api/sessions", s.listSessions)
|
||||||
api.HandleFunc("POST /api/sessions", s.createSession)
|
api.HandleFunc("POST /api/sessions", s.createSession)
|
||||||
api.HandleFunc("GET /api/sessions/{id}", s.getSession)
|
api.HandleFunc("GET /api/sessions/{id}", s.getSession)
|
||||||
|
|
@ -110,7 +121,7 @@ func (s *Server) originCheck(h http.Handler) http.Handler {
|
||||||
w.Header().Set("Vary", "Origin")
|
w.Header().Set("Vary", "Origin")
|
||||||
if r.Method == http.MethodOptions {
|
if r.Method == http.MethodOptions {
|
||||||
w.Header().Set("Access-Control-Allow-Headers", "Authorization, Content-Type")
|
w.Header().Set("Access-Control-Allow-Headers", "Authorization, Content-Type")
|
||||||
w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PATCH, DELETE")
|
w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, PATCH, DELETE")
|
||||||
w.WriteHeader(http.StatusNoContent)
|
w.WriteHeader(http.StatusNoContent)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
@ -252,6 +263,10 @@ func (s *Server) chat(w http.ResponseWriter, r *http.Request) {
|
||||||
ctx, cancel := context.WithTimeout(r.Context(), s.RunTimeout)
|
ctx, cancel := context.WithTimeout(r.Context(), s.RunTimeout)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
reply, err := s.Agent.Run(ctx, id, req.Text, nil)
|
reply, err := s.Agent.Run(ctx, id, req.Text, nil)
|
||||||
|
if errors.Is(err, agent.ErrCancelled) {
|
||||||
|
writeJSON(w, 200, map[string]any{"session_id": id, "reply": reply, "cancelled": true})
|
||||||
|
return
|
||||||
|
}
|
||||||
if err != nil {
|
if err != nil {
|
||||||
s.Log.Warn("agent run", "err", err)
|
s.Log.Warn("agent run", "err", err)
|
||||||
writeJSON(w, http.StatusBadGateway, map[string]string{"session_id": id, "error": err.Error()})
|
writeJSON(w, http.StatusBadGateway, map[string]string{"session_id": id, "error": err.Error()})
|
||||||
|
|
@ -260,8 +275,10 @@ func (s *Server) chat(w http.ResponseWriter, r *http.Request) {
|
||||||
writeJSON(w, 200, map[string]string{"session_id": id, "reply": reply})
|
writeJSON(w, 200, map[string]string{"session_id": id, "reply": reply})
|
||||||
}
|
}
|
||||||
|
|
||||||
// ws protocol: client sends {"type":"message","session_id":"...","text":"..."}
|
// ws protocol: client sends {"type":"message","session_id":"...","text":"..."};
|
||||||
// server streams agent.Event JSON objects ending with type "done" or "error".
|
// the server streams agent.Event JSON objects ending with "done", "cancelled"
|
||||||
|
// or "error". While an answer is streaming the client may send
|
||||||
|
// {"type":"cancel","session_id":"..."} (or call POST /api/sessions/{id}/cancel).
|
||||||
func (s *Server) ws(w http.ResponseWriter, r *http.Request) {
|
func (s *Server) ws(w http.ResponseWriter, r *http.Request) {
|
||||||
// Origin was already validated by originCheck; the library's own check is
|
// Origin was already validated by originCheck; the library's own check is
|
||||||
// therefore disabled here to avoid duplicating (and diverging from) policy.
|
// therefore disabled here to avoid duplicating (and diverging from) policy.
|
||||||
|
|
@ -271,37 +288,71 @@ func (s *Server) ws(w http.ResponseWriter, r *http.Request) {
|
||||||
}
|
}
|
||||||
defer c.CloseNow()
|
defer c.CloseNow()
|
||||||
c.SetReadLimit(maxBody)
|
c.SetReadLimit(maxBody)
|
||||||
ctx := r.Context()
|
ctx, stop := context.WithCancel(r.Context())
|
||||||
|
defer stop()
|
||||||
|
var wmu sync.Mutex
|
||||||
|
send := func(e agent.Event) {
|
||||||
|
wmu.Lock()
|
||||||
|
defer wmu.Unlock()
|
||||||
|
wctx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
wsjson.Write(wctx, c, e)
|
||||||
|
}
|
||||||
|
var running sync.WaitGroup
|
||||||
|
defer running.Wait()
|
||||||
|
busy := make(chan struct{}, 1) // one answer at a time per connection
|
||||||
for {
|
for {
|
||||||
var req chatReq
|
var req chatReq
|
||||||
if err := wsjson.Read(ctx, c, &req); err != nil {
|
if err := wsjson.Read(ctx, c, &req); err != nil {
|
||||||
|
stop() // connection gone: abort any running answer
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
if req.Type == "ping" {
|
switch req.Type {
|
||||||
wsjson.Write(ctx, c, agent.Event{Type: "pong"})
|
case "ping":
|
||||||
|
send(agent.Event{Type: "pong"})
|
||||||
|
continue
|
||||||
|
case "cancel":
|
||||||
|
if req.SessionID == "" || !s.Agent.Cancel(req.SessionID) {
|
||||||
|
send(agent.Event{Type: "error", Text: "nothing to cancel", Session: req.SessionID})
|
||||||
|
}
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
if strings.TrimSpace(req.Text) == "" {
|
if strings.TrimSpace(req.Text) == "" {
|
||||||
wsjson.Write(ctx, c, agent.Event{Type: "error", Text: "text required"})
|
send(agent.Event{Type: "error", Text: "text required"})
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case busy <- struct{}{}:
|
||||||
|
default:
|
||||||
|
send(agent.Event{Type: "error", Text: "an answer is still streaming; cancel it first"})
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
id, err := s.ensureSession(ctx, req.SessionID, req.Text)
|
id, err := s.ensureSession(ctx, req.SessionID, req.Text)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
wsjson.Write(ctx, c, agent.Event{Type: "error", Text: "unknown session"})
|
<-busy
|
||||||
|
send(agent.Event{Type: "error", Text: "unknown session"})
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
wsjson.Write(ctx, c, agent.Event{Type: "session", Session: id})
|
send(agent.Event{Type: "session", Session: id})
|
||||||
|
running.Add(1)
|
||||||
|
go func(id, text string) {
|
||||||
|
defer running.Done()
|
||||||
|
defer func() { <-busy }()
|
||||||
runCtx, cancel := context.WithTimeout(ctx, s.RunTimeout)
|
runCtx, cancel := context.WithTimeout(ctx, s.RunTimeout)
|
||||||
reply, err := s.Agent.Run(runCtx, id, req.Text, func(e agent.Event) {
|
defer cancel()
|
||||||
|
reply, err := s.Agent.Run(runCtx, id, text, func(e agent.Event) {
|
||||||
e.Session = id
|
e.Session = id
|
||||||
wsjson.Write(runCtx, c, e)
|
send(e)
|
||||||
})
|
})
|
||||||
cancel()
|
switch {
|
||||||
if err != nil {
|
case errors.Is(err, agent.ErrCancelled):
|
||||||
|
send(agent.Event{Type: "cancelled", Text: reply, Session: id})
|
||||||
|
case err != nil:
|
||||||
s.Log.Warn("agent run", "err", err)
|
s.Log.Warn("agent run", "err", err)
|
||||||
wsjson.Write(ctx, c, agent.Event{Type: "error", Text: err.Error(), Session: id})
|
send(agent.Event{Type: "error", Text: err.Error(), Session: id})
|
||||||
continue
|
default:
|
||||||
|
send(agent.Event{Type: "done", Text: reply, Session: id})
|
||||||
}
|
}
|
||||||
wsjson.Write(ctx, c, agent.Event{Type: "done", Text: reply, Session: id})
|
}(id, req.Text)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue