feat(agent): 支持持久化 /btw 旁路提问
This commit is contained in:
ji ruo committed
1 parent
8dae851b9b
commit
555539bc8e
34 files changed
+3744
-314
No files matched your search
@@ -16,6 +16,7 @@ import (
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/intercept"
|
||||
"github.com/Autumn-27/artex/sidequestion"
|
||||
)
|
||||
|
||||
const (
|
||||
@@ -204,6 +205,7 @@ func (s *Server) pgDeleteConversation(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
s.cancelConversation(c.ID)
|
||||
s.cancelSideWhere(func(p sidequestion.Parent) bool { return p.ConversationID == c.ID })
|
||||
if err := pg.DeleteConversation(c.ID); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
@@ -262,6 +264,7 @@ func (s *Server) pgDeleteConversationsBatch(w http.ResponseWriter, r *http.Reque
|
||||
}
|
||||
for _, id := range ids {
|
||||
s.cancelConversation(id)
|
||||
s.cancelSideWhere(func(p sidequestion.Parent) bool { return p.ConversationID == id })
|
||||
}
|
||||
deleted, err := pg.DeleteConversations(ids)
|
||||
if err != nil {
|
||||
|
||||
+27
-11
@@ -55,10 +55,11 @@ type Server struct {
|
||||
chatAgent *agent.ChatAgent // conversational runner for the chat page; nil w/o LLM
|
||||
// llmProv is the fully-decorated global provider (recorder + failover chain)
|
||||
// installed by applyLLM. Task routers use it after an explicit chain is cleared.
|
||||
llmProv llm.Provider
|
||||
llmCfg agent.Config // current LLM config (key not exposed)
|
||||
llmOn bool
|
||||
llmProf string // active LLM profile name (for llmrec tagging)
|
||||
llmProv llm.Provider
|
||||
llmDirect llm.Provider // concrete global provider, before any failover pool
|
||||
llmCfg agent.Config // current LLM config (key not exposed)
|
||||
llmOn bool
|
||||
llmProf string // active LLM profile name (for llmrec tagging)
|
||||
|
||||
// chatBusy guards the per-task main-agent run: the chat handler launches the
|
||||
// agent on the server's background ctx (not the request ctx) and returns
|
||||
@@ -121,6 +122,7 @@ type Server struct {
|
||||
// channel coalesces enqueue bursts; the database remains the source of truth.
|
||||
archiveWake chan struct{}
|
||||
archiveWG sync.WaitGroup
|
||||
side *sideQuestionState
|
||||
}
|
||||
|
||||
// profBundle is a planner/worker pair built from one LLM profile.
|
||||
@@ -157,6 +159,7 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi
|
||||
profAgents: map[int64]*profBundle{}, profChatAgents: map[int64]*agent.ChatAgent{},
|
||||
provByProfile: map[int64]*provEntry{}, llmHealth: newLLMHealthRegistry(m.pg),
|
||||
taskAgents: map[string]*taskAgentBundle{}, archiveWake: make(chan struct{}, 1)}
|
||||
s.initSideQuestions()
|
||||
// 熔断阈值/冷却是失败路径上的热参数，启动时把全局重试策略推给 Registry 一次；
|
||||
// 之后每次保存策略再推一次（saveLLMRetryPolicy）。
|
||||
s.applyRetryPolicy()
|
||||
@@ -243,6 +246,17 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi
|
||||
} else {
|
||||
log.Printf("[engine] no LLM provider configured — engine idle until set via /api/llm or env")
|
||||
}
|
||||
s.restoreTaskRuntimes()
|
||||
go s.reconcileConcurrency()
|
||||
s.startTaskArchiveWorker()
|
||||
s.wireInterceptReviewer() // LLM 兜底审批:未命中拦截规则的命令交给模型判定
|
||||
return s
|
||||
}
|
||||
|
||||
// Restored deadline and worker loops must inherit the same context as new tasks,
|
||||
// including the side-question checkpoint publisher installed during startup.
|
||||
func (s *Server) restoreTaskRuntimes() {
|
||||
m := s.m
|
||||
// reload tasks persisted on disk so the task list survives a restart, and
|
||||
// restore persisted paused state (so a task paused before restart stays paused).
|
||||
for _, t := range m.LoadExisting() {
|
||||
@@ -258,7 +272,7 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi
|
||||
// 任务级超时:为每个未终态、带 timeout 的任务起 deadline 协调器,独立于 planner/worker
|
||||
// loop——非活跃任务重启后也能在到点后被收尾(deadline 已过则立即走收尾时序)。
|
||||
if !isTerminalStatus(lifecycle.Status) {
|
||||
s.engine.startDeadlineCoordinator(ctx, t)
|
||||
s.engine.startDeadlineCoordinator(s.ctx, t)
|
||||
}
|
||||
}
|
||||
// Restore every task that had already been admitted before shutdown. Starting
|
||||
@@ -268,13 +282,9 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi
|
||||
for _, t := range m.List() {
|
||||
lifecycle := t.lifecycleSnapshot()
|
||||
if !lifecycle.Queued && !isTerminalStatus(lifecycle.Status) {
|
||||
s.engine.Run(ctx, t)
|
||||
s.engine.Run(s.ctx, t)
|
||||
}
|
||||
}
|
||||
go s.reconcileConcurrency()
|
||||
s.startTaskArchiveWorker()
|
||||
s.wireInterceptReviewer() // LLM 兜底审批:未命中拦截规则的命令交给模型判定
|
||||
return s
|
||||
}
|
||||
|
||||
// agentMaxTurns returns the configured max_turns for an agent key (0 = unlimited,
|
||||
@@ -360,7 +370,7 @@ func (s *Server) saveLLMConfig(cfg agent.Config) error {
|
||||
ContextWindowK: cfg.ContextWindowK, ThinkingType: cfg.ThinkingType, ReasoningEffort: cfg.ReasoningEffort, IsDefault: true,
|
||||
Priority: priority, PoolExclude: poolExclude, Streaming: streaming,
|
||||
MaxTokens: maxTokens, MaxTokensField: maxTokensField, SessionHeaderKey: sessionHeaderKey,
|
||||
Retry: retry,
|
||||
Retry: retry,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -464,7 +474,11 @@ func (s *Server) applyLLM(cfg agent.Config) error {
|
||||
profName := s.llmProf
|
||||
s.cfgMu.Unlock()
|
||||
prov = llmrec.Wrap(prov, s.m.PG(), cfg.Model, profName, cfg.ThinkingType, cfg.ReasoningEffort, s.m.LLMRecordEnabled)
|
||||
prov = bindSideProvider(prov, cfg, 0, profName)
|
||||
}
|
||||
s.cfgMu.Lock()
|
||||
s.llmDirect = prov
|
||||
s.cfgMu.Unlock()
|
||||
// LLM 轮询(默认关):把激活配置包进故障转移链,当前配置不可用时自动切下一个。
|
||||
// 只影响「走全局激活配置」的这条路径——agent 绑定 / 任务 pin 的走 providerForProfile,
|
||||
// 默认仍然独占该配置(见 poolForBinding)。关闭或无备选时返回原 provider,行为不变。
|
||||
@@ -598,6 +612,7 @@ func (s *Server) providerForProfile(id int64) (llm.Provider, agent.Config, bool)
|
||||
// Wrap with recorder, tagged with this profile's name.
|
||||
if p, _ := s.m.pg.ProfileByID(id); p != nil {
|
||||
prov = llmrec.Wrap(prov, s.m.PG(), cfg.Model, p.Name, cfg.ThinkingType, cfg.ReasoningEffort, s.m.LLMRecordEnabled)
|
||||
prov = bindSideProvider(prov, cfg, id, p.Name)
|
||||
}
|
||||
s.provCacheMu.Lock()
|
||||
if generation != s.provCacheGen {
|
||||
@@ -716,6 +731,7 @@ func (s *Server) chatAgentRef() *agent.ChatAgent {
|
||||
|
||||
func (s *Server) Handler() http.Handler {
|
||||
mux := http.NewServeMux()
|
||||
s.registerSideRoutes(mux)
|
||||
|
||||
// Auth routes — exempt from JWT check (handled in requireAuth)
|
||||
mux.HandleFunc("GET /api/auth/status", s.authStatus)
|
||||
|
||||
@@ -162,6 +162,10 @@ func (s *Server) pgDeleteTask(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
if err := s.drainTaskSideQuestions(drainCtx, id); err != nil {
|
||||
writeErr(w, http.StatusConflict, err.Error())
|
||||
return
|
||||
}
|
||||
result, err := s.m.DeleteTask(id, opts)
|
||||
if err != nil {
|
||||
var committed *taskDeleteCommittedError
|
||||
|
||||
@@ -0,0 +1,574 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/llmrec"
|
||||
"github.com/Autumn-27/artex/sidequestion"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
"github.com/Autumn-27/norma/transcript"
|
||||
)
|
||||
|
||||
type sideRun struct {
|
||||
key string
|
||||
parent sidequestion.Parent
|
||||
cancel context.CancelFunc
|
||||
done chan struct{}
|
||||
}
|
||||
type sideQuestionState struct {
|
||||
mu sync.Mutex
|
||||
commands sync.Mutex // short admission / clear operations only, never inference
|
||||
pending map[string]sidequestion.Snapshot
|
||||
latest map[string]sidequestion.Snapshot
|
||||
seen map[string][2]int64
|
||||
runs map[string]sideRun
|
||||
done chan struct{}
|
||||
}
|
||||
|
||||
func sideModel(cfg agent.Config, id int64, name string) sidequestion.Model {
|
||||
b, _ := json.Marshal([]any{cfg.Format, cfg.BaseURL, cfg.Model, cfg.ThinkingType, cfg.ReasoningEffort, cfg.Stream, cfg.MaxTokensField, cfg.SessionHeaderKey})
|
||||
h := sha256.Sum256(b)
|
||||
return sidequestion.Model{ProfileID: id, Name: name, Format: string(cfg.Format), Model: cfg.Model, Identity: hex.EncodeToString(h[:]), Streaming: cfg.Stream, WindowTokens: cfg.CompactionWindow()}
|
||||
}
|
||||
|
||||
func bindSideProvider(p llm.Provider, cfg agent.Config, id int64, name string) llm.Provider {
|
||||
return sidequestion.Bind(p, sideModel(cfg, id, name))
|
||||
}
|
||||
|
||||
func (s *Server) initSideQuestions() {
|
||||
s.side = &sideQuestionState{pending: map[string]sidequestion.Snapshot{}, latest: map[string]sidequestion.Snapshot{}, seen: map[string][2]int64{}, runs: map[string]sideRun{}, done: make(chan struct{})}
|
||||
if s.m.pg == nil {
|
||||
close(s.side.done)
|
||||
return
|
||||
}
|
||||
if err := s.m.pg.InterruptSideRequests(s.ctx); err != nil {
|
||||
log.Printf("[btw] recover: %v", err)
|
||||
}
|
||||
s.ctx = sidequestion.WithPublisher(s.ctx, func(snap sidequestion.Snapshot) {
|
||||
s.side.mu.Lock()
|
||||
defer s.side.mu.Unlock()
|
||||
key := snap.Parent.Key()
|
||||
old, ok := s.side.seen[key]
|
||||
if ok && (old[0] > snap.RunID || old[0] == snap.RunID && old[1] >= snap.Version) {
|
||||
return
|
||||
}
|
||||
s.side.latest[key] = snap
|
||||
s.side.seen[key] = [2]int64{snap.RunID, snap.Version}
|
||||
s.side.pending[key] = snap
|
||||
})
|
||||
go func() {
|
||||
defer close(s.side.done)
|
||||
tick := time.NewTicker(250 * time.Millisecond)
|
||||
defer tick.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-tick.C:
|
||||
s.flushSideSnapshots()
|
||||
case <-s.ctx.Done():
|
||||
s.flushSideSnapshots()
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func (s *Server) flushSideSnapshots() {
|
||||
if s.side == nil || s.m.pg == nil {
|
||||
return
|
||||
}
|
||||
s.side.commands.Lock()
|
||||
defer s.side.commands.Unlock()
|
||||
s.side.mu.Lock()
|
||||
pending := s.side.pending
|
||||
s.side.pending = map[string]sidequestion.Snapshot{}
|
||||
s.side.mu.Unlock()
|
||||
for key, snap := range pending {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
err := s.m.pg.SaveSideSnapshot(ctx, snap)
|
||||
cancel()
|
||||
s.side.mu.Lock()
|
||||
if err == nil || errors.Is(err, db.ErrSideParentGone) {
|
||||
if cur, ok := s.side.latest[key]; ok && cur.RunID == snap.RunID && cur.Version == snap.Version {
|
||||
delete(s.side.latest, key)
|
||||
}
|
||||
} else if newer, ok := s.side.pending[key]; !ok || newer.RunID < snap.RunID || newer.RunID == snap.RunID && newer.Version < snap.Version {
|
||||
s.side.pending[key] = snap
|
||||
}
|
||||
s.side.mu.Unlock()
|
||||
if err != nil && !errors.Is(err, db.ErrSideParentGone) {
|
||||
log.Printf("[btw] checkpoint persistence: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) sideSnapshot(ctx context.Context, p sidequestion.Parent) (*sidequestion.Snapshot, error) {
|
||||
s.side.mu.Lock()
|
||||
snap, ok := s.side.latest[p.Key()]
|
||||
s.side.mu.Unlock()
|
||||
if ok {
|
||||
return &snap, nil
|
||||
}
|
||||
return s.m.pg.SideSnapshot(ctx, p.Key())
|
||||
}
|
||||
|
||||
func (s *Server) cancelSideWhere(match func(sidequestion.Parent) bool) []<-chan struct{} {
|
||||
if s.side == nil {
|
||||
return nil
|
||||
}
|
||||
s.side.mu.Lock()
|
||||
defer s.side.mu.Unlock()
|
||||
var done []<-chan struct{}
|
||||
for _, run := range s.side.runs {
|
||||
if match(run.parent) {
|
||||
run.cancel()
|
||||
done = append(done, run.done)
|
||||
}
|
||||
}
|
||||
return done
|
||||
}
|
||||
|
||||
// Called after the task admission barrier closes and the main loops drain.
|
||||
// Wait for final answer/usage writes before taking the archive snapshot.
|
||||
func (s *Server) drainTaskSideQuestions(ctx context.Context, taskID string) error {
|
||||
if s.side == nil {
|
||||
return nil
|
||||
}
|
||||
id, _ := strconv.ParseInt(taskID, 10, 64)
|
||||
s.side.commands.Lock()
|
||||
done := s.cancelSideWhere(func(p sidequestion.Parent) bool { return p.TaskID == id })
|
||||
s.side.commands.Unlock()
|
||||
for _, ch := range done {
|
||||
select {
|
||||
case <-ch:
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
s.flushSideSnapshots()
|
||||
// A failed flush must not silently leave the latest checkpoint out of archive.
|
||||
s.side.mu.Lock()
|
||||
defer s.side.mu.Unlock()
|
||||
for _, snap := range s.side.pending {
|
||||
if snap.Parent.TaskID == id {
|
||||
return errors.New("旁路上下文尚未保存，请重试")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Server) cancelWorkerSide(taskID string, expID, iid int64) {
|
||||
if s.side == nil {
|
||||
return
|
||||
}
|
||||
id, _ := strconv.ParseInt(taskID, 10, 64)
|
||||
p := sidequestion.Parent{TaskID: id, ExplorationID: expID, IntentID: iid}
|
||||
s.side.commands.Lock()
|
||||
defer s.side.commands.Unlock()
|
||||
s.cancelSideWhere(func(parent sidequestion.Parent) bool { return parent == p })
|
||||
}
|
||||
|
||||
func (s *Server) sideProvider(model sidequestion.Model) (llm.Provider, error) {
|
||||
var p llm.Provider
|
||||
var cfg agent.Config
|
||||
var ok bool
|
||||
if model.ProfileID > 0 {
|
||||
// Validate the persisted reference even if a previous provider is cached.
|
||||
current, exists := s.loadProfileConfig(model.ProfileID)
|
||||
if !exists || sideModel(current, model.ProfileID, model.Name).Identity != model.Identity {
|
||||
return nil, errors.New("模型配置已删除或变化，请先运行主 Agent 更新上下文")
|
||||
}
|
||||
p, cfg, ok = s.providerForProfile(model.ProfileID)
|
||||
} else {
|
||||
s.cfgMu.Lock()
|
||||
p, cfg, ok = s.llmDirect, s.llmCfg, s.llmOn
|
||||
s.cfgMu.Unlock()
|
||||
}
|
||||
if !ok || p == nil || sideModel(cfg, model.ProfileID, model.Name).Identity != model.Identity {
|
||||
return nil, errors.New("模型配置已删除或变化，请先运行主 Agent 更新上下文")
|
||||
}
|
||||
return p, nil
|
||||
}
|
||||
|
||||
func (s *Server) sideParent(w http.ResponseWriter, r *http.Request, kind string) (sidequestion.Parent, bool) {
|
||||
p := sidequestion.Parent{}
|
||||
if s.side == nil || s.m.pg == nil {
|
||||
writeErr(w, 503, "旁路服务不可用")
|
||||
return p, false
|
||||
}
|
||||
id, err := strconv.ParseInt(r.PathValue("id"), 10, 64)
|
||||
if err != nil || id <= 0 {
|
||||
writeErr(w, 400, "bad id")
|
||||
return p, false
|
||||
}
|
||||
if kind == "conversation" {
|
||||
c, err := s.m.pg.GetConversation(id)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return p, false
|
||||
}
|
||||
if c == nil {
|
||||
writeErr(w, 404, "conversation not found")
|
||||
return p, false
|
||||
}
|
||||
p.ConversationID = id
|
||||
return p, true
|
||||
}
|
||||
t, ok := s.m.Task(r.PathValue("id"))
|
||||
if !ok {
|
||||
writeErr(w, 404, "task not found")
|
||||
return p, false
|
||||
}
|
||||
pt, err := s.m.pg.GetTask(id)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return p, false
|
||||
}
|
||||
if pt == nil || s.engine.IsDeleting(t.ID) {
|
||||
writeErr(w, 409, "任务已归档或正在删除")
|
||||
return p, false
|
||||
}
|
||||
p.TaskID, p.ExplorationID = id, t.ExpID
|
||||
if kind == "worker" {
|
||||
iid, err := strconv.ParseInt(r.PathValue("iid"), 10, 64)
|
||||
if err != nil || iid <= 0 {
|
||||
writeErr(w, 400, "bad intent id")
|
||||
return p, false
|
||||
}
|
||||
n, err := t.Store.GetNode(iid)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return p, false
|
||||
}
|
||||
if n == nil || n.Kind != db.KindIntent {
|
||||
writeErr(w, 404, "intent not found")
|
||||
return p, false
|
||||
}
|
||||
if n.State == "stopped" {
|
||||
writeErr(w, 409, "Worker 已删除")
|
||||
return p, false
|
||||
}
|
||||
p.IntentID = iid
|
||||
}
|
||||
return p, true
|
||||
}
|
||||
|
||||
func (s *Server) registerSideRoutes(mux *http.ServeMux) {
|
||||
for _, route := range []struct{ path, kind string }{{"/api/conversations/{id}", "conversation"}, {"/api/tasks/{id}/chat", "main"}, {"/api/tasks/{id}/intents/{iid}", "worker"}} {
|
||||
for _, method := range []string{"GET", "POST", "DELETE"} {
|
||||
mux.HandleFunc(method+" "+route.path+"/side-questions", func(w http.ResponseWriter, r *http.Request) {
|
||||
p, ok := s.sideParent(w, r, route.kind)
|
||||
if ok {
|
||||
s.handleSideQuestions(w, r, p)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
mux.HandleFunc("GET /api/side-questions/{requestID}/events", s.sideEvents)
|
||||
mux.HandleFunc("POST /api/side-questions/{requestID}/cancel", s.cancelSideRequest)
|
||||
}
|
||||
|
||||
func (s *Server) handleSideQuestions(w http.ResponseWriter, r *http.Request, p sidequestion.Parent) {
|
||||
key := p.Key()
|
||||
if r.Method == "DELETE" {
|
||||
s.side.commands.Lock()
|
||||
defer s.side.commands.Unlock()
|
||||
s.side.mu.Lock()
|
||||
for _, run := range s.side.runs {
|
||||
if run.key == key {
|
||||
run.cancel()
|
||||
}
|
||||
}
|
||||
s.side.mu.Unlock()
|
||||
if err := s.m.pg.ClearSideHistory(r.Context(), key); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, map[string]bool{"cleared": true})
|
||||
return
|
||||
}
|
||||
snap, err := s.sideSnapshot(r.Context(), p)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if r.Method == "GET" {
|
||||
before, _ := strconv.ParseInt(r.URL.Query().Get("before"), 10, 64)
|
||||
items, err := s.m.pg.SideHistory(r.Context(), key, before, 21)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
var next int64
|
||||
if len(items) > 20 {
|
||||
items = items[:20]
|
||||
next = items[len(items)-1].Ordinal
|
||||
}
|
||||
current, err := s.m.pg.CurrentSideRequest(r.Context(), key)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
var meta any
|
||||
if snap != nil {
|
||||
_, modelErr := s.sideProvider(snap.Model)
|
||||
reason := ""
|
||||
if modelErr != nil {
|
||||
reason = modelErr.Error()
|
||||
}
|
||||
meta = map[string]any{"captured_at": snap.CapturedAt, "model": snap.Model, "available": modelErr == nil, "reason": reason}
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"items": items, "current": current, "next_cursor": next, "snapshot": meta})
|
||||
return
|
||||
}
|
||||
r.Body = http.MaxBytesReader(w, r.Body, 32<<10)
|
||||
var in struct {
|
||||
Question string `json:"question"`
|
||||
ClientID string `json:"client_request_id"`
|
||||
}
|
||||
if json.NewDecoder(r.Body).Decode(&in) != nil {
|
||||
writeErr(w, 400, "bad json")
|
||||
return
|
||||
}
|
||||
in.Question = strings.TrimSpace(in.Question)
|
||||
if in.Question == "" || len([]rune(in.Question)) > 4000 || !validWorkerMessageRequestID(in.ClientID) {
|
||||
writeErr(w, 400, "问题须为 1–4000 字符，并提供有效请求 ID")
|
||||
return
|
||||
}
|
||||
s.side.commands.Lock()
|
||||
defer s.side.commands.Unlock()
|
||||
if p.TaskID > 0 && s.engine.IsDeleting(strconv.FormatInt(p.TaskID, 10)) {
|
||||
writeErr(w, 409, "任务正在归档或删除")
|
||||
return
|
||||
}
|
||||
if existing, err := s.m.pg.ExistingSideRequest(r.Context(), key, in.ClientID); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
} else if existing != nil {
|
||||
if existing.Question != in.Question {
|
||||
writeErr(w, 409, "同一请求 ID 不能用于不同问题")
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, existing)
|
||||
return
|
||||
}
|
||||
if snap == nil {
|
||||
writeErr(w, 409, "尚无上下文快照，请先运行主 Agent")
|
||||
return
|
||||
}
|
||||
provider, err := s.sideProvider(snap.Model)
|
||||
if err != nil {
|
||||
writeErr(w, 409, err.Error())
|
||||
return
|
||||
}
|
||||
s.side.mu.Lock()
|
||||
full := len(s.side.runs) >= 4
|
||||
busy := false
|
||||
for _, run := range s.side.runs {
|
||||
if run.key == key {
|
||||
busy = true
|
||||
}
|
||||
}
|
||||
s.side.mu.Unlock()
|
||||
if busy {
|
||||
writeErr(w, 409, db.ErrSideBusy.Error())
|
||||
return
|
||||
}
|
||||
if full {
|
||||
writeErr(w, 429, "旁路请求已达并发上限，请稍后重试")
|
||||
return
|
||||
}
|
||||
if err = s.m.pg.SaveSideSnapshot(r.Context(), *snap); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
history, err := s.m.pg.SideReplay(r.Context(), key)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
req, err := sidequestion.BuildRequest(*snap, history, in.Question)
|
||||
if err != nil {
|
||||
writeErr(w, 409, err.Error())
|
||||
return
|
||||
}
|
||||
e, created, err := s.m.pg.StartSideRequest(r.Context(), *snap, in.ClientID, in.Question)
|
||||
if err != nil {
|
||||
writeErr(w, 409, err.Error())
|
||||
return
|
||||
}
|
||||
if created {
|
||||
ctx, cancel := context.WithTimeout(s.ctx, 120*time.Second)
|
||||
s.side.mu.Lock()
|
||||
s.side.runs[e.ID] = sideRun{key: key, parent: p, cancel: cancel, done: make(chan struct{})}
|
||||
s.side.mu.Unlock()
|
||||
go s.runSide(ctx, cancel, *e, p, provider, req)
|
||||
}
|
||||
writeJSON(w, 202, e)
|
||||
}
|
||||
|
||||
func (s *Server) runSide(ctx context.Context, cancel context.CancelFunc, e sidequestion.Exchange, parent sidequestion.Parent, provider llm.Provider, req llm.CompletionRequest) {
|
||||
defer cancel()
|
||||
defer func() {
|
||||
s.side.mu.Lock()
|
||||
if run, ok := s.side.runs[e.ID]; ok {
|
||||
close(run.done)
|
||||
}
|
||||
delete(s.side.runs, e.ID)
|
||||
s.side.mu.Unlock()
|
||||
}()
|
||||
ctx = transcript.WithSessionID(ctx, fmt.Sprintf("exp%d-btw-%s", parent.ExplorationID, e.ID))
|
||||
if parent.TaskID > 0 {
|
||||
ctx = llmrec.WithTaskID(ctx, strconv.FormatInt(parent.TaskID, 10))
|
||||
}
|
||||
// Deletions can race between request insertion and registration. Observe the
|
||||
// persisted lifecycle even while a provider is blocked without emitting text.
|
||||
go func() {
|
||||
tick := time.NewTicker(500 * time.Millisecond)
|
||||
defer tick.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-tick.C:
|
||||
row, err := s.m.pg.SideRequest(ctx, e.ID)
|
||||
if err == nil && (row == nil || !row.Running()) {
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
lastSave := time.Now()
|
||||
persist := func() (bool, error) {
|
||||
e.Sequence++
|
||||
writeCtx, stop := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
defer stop()
|
||||
ok, err := s.m.pg.UpdateSideRequest(writeCtx, e)
|
||||
if err != nil {
|
||||
log.Printf("[btw] answer persistence: %v", err)
|
||||
return true, err
|
||||
}
|
||||
return ok, nil
|
||||
}
|
||||
answer, runErr := (sidequestion.SideQuestionService{Provider: provider}).Answer(ctx, req, e.Model.Streaming, func(a sidequestion.Answer) {
|
||||
e.Answer, e.Usage = a.Text, a.Usage
|
||||
if time.Since(lastSave) >= 250*time.Millisecond {
|
||||
lastSave = time.Now()
|
||||
if ok, _ := persist(); !ok {
|
||||
cancel()
|
||||
}
|
||||
}
|
||||
})
|
||||
e.Answer, e.Usage = answer.Text, answer.Usage
|
||||
e.Status = "completed"
|
||||
if ctx.Err() != nil {
|
||||
e.Status = "cancelled"
|
||||
e.Error = "回答已停止"
|
||||
if errors.Is(ctx.Err(), context.DeadlineExceeded) {
|
||||
e.Status = "failed"
|
||||
e.Error = "旁路回答超过 120 秒，已停止"
|
||||
}
|
||||
} else if runErr != nil {
|
||||
e.Status = "failed"
|
||||
e.Error = runErr.Error()
|
||||
}
|
||||
for attempt := 0; attempt < 3; attempt++ {
|
||||
if _, err := persist(); err == nil {
|
||||
break
|
||||
}
|
||||
if attempt < 2 {
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) cancelSideRequest(w http.ResponseWriter, r *http.Request) {
|
||||
e, err := s.m.pg.SideRequest(r.Context(), r.PathValue("requestID"))
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if e == nil {
|
||||
writeErr(w, 404, "side question not found")
|
||||
return
|
||||
}
|
||||
s.side.mu.Lock()
|
||||
if run, ok := s.side.runs[e.ID]; ok {
|
||||
run.cancel()
|
||||
}
|
||||
s.side.mu.Unlock()
|
||||
writeJSON(w, 200, map[string]bool{"cancelled": true})
|
||||
}
|
||||
|
||||
func (s *Server) sideEvents(w http.ResponseWriter, r *http.Request) {
|
||||
flusher, ok := w.(http.Flusher)
|
||||
if !ok {
|
||||
writeErr(w, 500, "streaming unavailable")
|
||||
return
|
||||
}
|
||||
id := r.PathValue("requestID")
|
||||
first, err := s.m.pg.SideRequest(r.Context(), id)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if first == nil {
|
||||
writeErr(w, 404, "side question not found")
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Cache-Control", "no-cache")
|
||||
w.Header().Set("X-Accel-Buffering", "no")
|
||||
seq := int64(-1)
|
||||
tick := time.NewTicker(250 * time.Millisecond)
|
||||
defer tick.Stop()
|
||||
heartbeat := time.NewTicker(15 * time.Second)
|
||||
defer heartbeat.Stop()
|
||||
for {
|
||||
e := first
|
||||
first = nil
|
||||
if e == nil {
|
||||
e, err = s.m.pg.SideRequest(r.Context(), id)
|
||||
}
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if e == nil {
|
||||
fmt.Fprint(w, "event: cleared\ndata: {}\n\n")
|
||||
flusher.Flush()
|
||||
return
|
||||
}
|
||||
if e.Sequence > seq {
|
||||
b, _ := json.Marshal(e)
|
||||
fmt.Fprintf(w, "id: %d\nevent: snapshot\ndata: %s\n\n", e.Sequence, b)
|
||||
flusher.Flush()
|
||||
seq = e.Sequence
|
||||
}
|
||||
if !e.Running() {
|
||||
return
|
||||
}
|
||||
select {
|
||||
case <-r.Context().Done():
|
||||
return
|
||||
case <-s.ctx.Done():
|
||||
return
|
||||
case <-tick.C:
|
||||
case <-heartbeat.C:
|
||||
fmt.Fprint(w, ": keepalive\n\n")
|
||||
flusher.Flush()
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,453 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"iter"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/sidequestion"
|
||||
"github.com/Autumn-27/norma/harness"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
)
|
||||
|
||||
type sideHTTPProvider struct {
|
||||
started chan llm.CompletionRequest
|
||||
release chan struct{}
|
||||
}
|
||||
|
||||
func (p *sideHTTPProvider) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
|
||||
return func(y func(llm.StreamEvent, error) bool) {
|
||||
p.started <- req
|
||||
if !y(llm.StreamEvent{Type: llm.SEMessageStart, Usage: llm.Usage{InputTokens: 19}}, nil) {
|
||||
return
|
||||
}
|
||||
if !y(llm.StreamEvent{Type: llm.SETextDelta, Text: "partial answer"}, nil) {
|
||||
return
|
||||
}
|
||||
select {
|
||||
case <-p.release:
|
||||
y(llm.StreamEvent{Type: llm.SEMessageStop}, nil)
|
||||
case <-ctx.Done():
|
||||
y(llm.StreamEvent{}, ctx.Err())
|
||||
}
|
||||
}
|
||||
}
|
||||
func (p *sideHTTPProvider) Complete(ctx context.Context, req llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
|
||||
for _, err := range p.Stream(ctx, req) {
|
||||
if err != nil {
|
||||
return llm.Message{}, "", llm.Usage{}, err
|
||||
}
|
||||
}
|
||||
return llm.Message{Role: llm.RoleAssistant, Content: []llm.ContentBlock{llm.TextBlock("atomic answer")}}, "end_turn", llm.Usage{InputTokens: 19}, nil
|
||||
}
|
||||
|
||||
type sideHTTPFixture struct {
|
||||
s *Server
|
||||
m *Manager
|
||||
provider *sideHTTPProvider
|
||||
handler http.Handler
|
||||
token string
|
||||
}
|
||||
|
||||
func newSideHTTPFixture(t *testing.T) *sideHTTPFixture {
|
||||
t.Helper()
|
||||
m, err := NewManager(t.TempDir(), "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
s := &Server{m: m, engine: NewEngine(m), ctx: ctx, jwtKey: []byte("btw-test-signing-key-only"), chatBusy: map[string]bool{}}
|
||||
s.initSideQuestions()
|
||||
p := &sideHTTPProvider{started: make(chan llm.CompletionRequest, 16), release: make(chan struct{})}
|
||||
cfg := agent.Config{Format: llm.FormatOpenAI, BaseURL: "http://fixture.invalid", Model: "fixture", Stream: true}
|
||||
s.cfgMu.Lock()
|
||||
s.llmCfg = cfg
|
||||
s.llmOn = true
|
||||
s.llmDirect = bindSideProvider(p, cfg, 0, "fixture")
|
||||
s.cfgMu.Unlock()
|
||||
token, err := signJWT(s.jwtKey)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
f := &sideHTTPFixture{s: s, m: m, provider: p, handler: s.Handler(), token: token}
|
||||
t.Cleanup(func() {
|
||||
cancel()
|
||||
for _, done := range s.cancelSideWhere(func(sidequestion.Parent) bool { return true }) {
|
||||
<-done
|
||||
}
|
||||
<-s.side.done
|
||||
s.flushSideSnapshots()
|
||||
m.Close()
|
||||
})
|
||||
return f
|
||||
}
|
||||
func (f *sideHTTPFixture) call(t *testing.T, method, path string, body any, want int) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
raw, _ := json.Marshal(body)
|
||||
req := httptest.NewRequest(method, path, bytes.NewReader(raw))
|
||||
req.Header.Set("Authorization", "Bearer "+f.token)
|
||||
w := httptest.NewRecorder()
|
||||
f.handler.ServeHTTP(w, req)
|
||||
if w.Code != want {
|
||||
t.Fatalf("%s %s: status=%d want=%d body=%s", method, path, w.Code, want, w.Body.String())
|
||||
}
|
||||
return w
|
||||
}
|
||||
func (f *sideHTTPFixture) conversation(t *testing.T) (sidequestion.Parent, string) {
|
||||
t.Helper()
|
||||
c, err := f.m.pg.CreateConversation("mainagent", "side HTTP", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = f.m.pg.DeleteConversation(c.ID) })
|
||||
return sidequestion.Parent{ConversationID: c.ID}, fmt.Sprintf("/api/conversations/%d/side-questions", c.ID)
|
||||
}
|
||||
func (f *sideHTTPFixture) checkpoint(t *testing.T, p sidequestion.Parent) sidequestion.Snapshot {
|
||||
t.Helper()
|
||||
snap := sidequestion.Snapshot{Parent: p, RunID: time.Now().UnixNano(), Version: 1, CapturedAt: time.Now().UTC(), Model: sideModel(f.s.llmCfg, 0, "fixture"), Request: llm.CompletionRequest{MaxTokens: 128, Messages: []llm.Message{llm.UserText("main context marker")}}}
|
||||
if err := f.m.pg.SaveSideSnapshot(t.Context(), snap); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return snap
|
||||
}
|
||||
func decodeSide(t *testing.T, w *httptest.ResponseRecorder) sidequestion.Exchange {
|
||||
t.Helper()
|
||||
var e sidequestion.Exchange
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &e); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return e
|
||||
}
|
||||
func waitSide(t *testing.T, d *db.DB, id, status string) *sidequestion.Exchange {
|
||||
t.Helper()
|
||||
until := time.Now().Add(3 * time.Second)
|
||||
for time.Now().Before(until) {
|
||||
e, err := d.SideRequest(t.Context(), id)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if e != nil && e.Status == status {
|
||||
return e
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
t.Fatalf("side %s did not reach %s", id, status)
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestSideHTTPBusyIsolationClearAndReconnect(t *testing.T) {
|
||||
f := newSideHTTPFixture(t)
|
||||
p, path := f.conversation(t)
|
||||
f.call(t, "POST", path, map[string]string{"question": "old session", "client_request_id": "old"}, 409)
|
||||
f.checkpoint(t, p)
|
||||
// Main busy status must not reject the side request.
|
||||
f.s.chatMu.Lock()
|
||||
f.s.chatBusy[fmt.Sprintf("conv-%d", p.ConversationID)] = true
|
||||
f.s.chatMu.Unlock()
|
||||
e := decodeSide(t, f.call(t, "POST", path, map[string]string{"question": "question", "client_request_id": "one"}, 202))
|
||||
select {
|
||||
case req := <-f.provider.started:
|
||||
if !strings.Contains(req.Messages[0].Text(), "main context marker") {
|
||||
t.Fatal("missing checkpoint")
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("side blocked by main busy")
|
||||
}
|
||||
again := decodeSide(t, f.call(t, "POST", path, map[string]string{"question": "question", "client_request_id": "one"}, 200))
|
||||
if again.ID != e.ID {
|
||||
t.Fatal("duplicate request re-executed")
|
||||
}
|
||||
f.call(t, "POST", path, map[string]string{"question": "another", "client_request_id": "two"}, 409)
|
||||
// SSE reconnect always starts with a cumulative snapshot even with Last-Event-ID.
|
||||
httpServer := httptest.NewServer(f.handler)
|
||||
defer httpServer.Close()
|
||||
for i := 0; i < 2; i++ {
|
||||
req, _ := http.NewRequest("GET", httpServer.URL+"/api/side-questions/"+e.ID+"/events", nil)
|
||||
req.Header.Set("Authorization", "Bearer "+f.token)
|
||||
req.Header.Set("Last-Event-ID", "999")
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
seen := false
|
||||
for scanner.Scan() {
|
||||
if strings.HasPrefix(scanner.Text(), "data: ") {
|
||||
seen = strings.Contains(scanner.Text(), e.ID)
|
||||
break
|
||||
}
|
||||
}
|
||||
resp.Body.Close()
|
||||
if !seen {
|
||||
t.Fatal("reconnect did not return cumulative request")
|
||||
}
|
||||
}
|
||||
f.call(t, "POST", "/api/side-questions/"+e.ID+"/cancel", nil, 200)
|
||||
done := waitSide(t, f.m.pg, e.ID, "cancelled")
|
||||
if done.Answer != "partial answer" || done.Usage.InputTokens != 19 {
|
||||
t.Fatalf("cancel lost partial/usage: %+v", done)
|
||||
}
|
||||
f.s.chatMu.Lock()
|
||||
mainBusy := f.s.chatBusy[fmt.Sprintf("conv-%d", p.ConversationID)]
|
||||
f.s.chatMu.Unlock()
|
||||
if !mainBusy {
|
||||
t.Fatal("side cancellation stopped main")
|
||||
}
|
||||
f.call(t, "DELETE", path, nil, 200)
|
||||
if history, err := f.m.pg.SideHistory(t.Context(), p.Key(), 0, 20); err != nil || len(history) != 0 {
|
||||
t.Fatalf("clear %+v %v", history, err)
|
||||
}
|
||||
if snap, err := f.m.pg.SideSnapshot(t.Context(), p.Key()); err != nil || snap == nil {
|
||||
t.Fatal("clear lost main snapshot")
|
||||
}
|
||||
var count int
|
||||
if err := f.m.pg.QueryRow(`SELECT count(*) FROM conversation_activities WHERE conversation_id=$1`, p.ConversationID).Scan(&count); err != nil || count != 0 {
|
||||
t.Fatalf("side contaminated main activity %d %v", count, err)
|
||||
}
|
||||
f.s.cfgMu.Lock()
|
||||
f.s.llmCfg.Model = "changed"
|
||||
f.s.cfgMu.Unlock()
|
||||
f.call(t, "GET", path, nil, 200)
|
||||
f.call(t, "POST", path, map[string]string{"question": "changed config", "client_request_id": "new"}, 409)
|
||||
}
|
||||
|
||||
func TestSideHTTPGlobalLimitTaskWorkerAndDeletion(t *testing.T) {
|
||||
f := newSideHTTPFixture(t)
|
||||
var requests []sidequestion.Exchange
|
||||
for i := 0; i < 5; i++ {
|
||||
p, path := f.conversation(t)
|
||||
f.checkpoint(t, p)
|
||||
want := 202
|
||||
if i == 4 {
|
||||
want = 429
|
||||
}
|
||||
w := f.call(t, "POST", path, map[string]string{"question": "blocked", "client_request_id": "id"}, want)
|
||||
if want == 202 {
|
||||
requests = append(requests, decodeSide(t, w))
|
||||
}
|
||||
}
|
||||
for _, e := range requests {
|
||||
f.call(t, "POST", "/api/side-questions/"+e.ID+"/cancel", nil, 200)
|
||||
waitSide(t, f.m.pg, e.ID, "cancelled")
|
||||
}
|
||||
close(f.provider.release)
|
||||
task, err := f.m.CreateTask("side route task", "context", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
id, _ := strconv.ParseInt(task.ID, 10, 64)
|
||||
t.Cleanup(func() { _ = f.m.pg.DeleteTask(id) })
|
||||
iid, err := task.Store.AddNode(db.KindIntent, map[string]any{"summary": "own worker"}, 1, "paused", "planner", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, intent := range []int64{0, iid} {
|
||||
p := sidequestion.Parent{TaskID: id, ExplorationID: task.ExpID, IntentID: intent}
|
||||
f.checkpoint(t, p)
|
||||
path := "/api/tasks/" + task.ID + "/chat/side-questions"
|
||||
if intent > 0 {
|
||||
path = fmt.Sprintf("/api/tasks/%s/intents/%d/side-questions", task.ID, intent)
|
||||
}
|
||||
e := decodeSide(t, f.call(t, "POST", path, map[string]string{"question": "which context", "client_request_id": "own"}, 202))
|
||||
waitSide(t, f.m.pg, e.ID, "completed")
|
||||
}
|
||||
f.call(t, "GET", fmt.Sprintf("/api/tasks/%s/intents/%d/side-questions", task.ID, iid+1000000), nil, 404)
|
||||
if _, err := f.s.applyIntentControl(t.Context(), task, iid, "cancel", "cleanup"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if snap, err := f.m.pg.SideSnapshot(t.Context(), (sidequestion.Parent{TaskID: id, ExplorationID: task.ExpID, IntentID: iid}).Key()); err != nil || snap != nil {
|
||||
t.Fatal("stopped intent side snapshot remains")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSideCheckpointPersistsBeforeAdmissionAndRestart(t *testing.T) {
|
||||
f := newSideHTTPFixture(t)
|
||||
p, path := f.conversation(t)
|
||||
bound := bindSideProvider(f.provider, f.s.llmCfg, 0, "fixture")
|
||||
ctx, deps := sidequestion.Attach(f.s.ctx, p, harness.QueryDeps{}, bound)
|
||||
mainCtx, stopMain := context.WithCancel(ctx)
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
for range deps.CallModel(mainCtx, llm.CompletionRequest{Messages: []llm.Message{llm.UserText("live checkpoint")}, MaxTokens: 128}) {
|
||||
}
|
||||
}()
|
||||
select {
|
||||
case <-f.provider.started:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("main never started")
|
||||
}
|
||||
e := decodeSide(t, f.call(t, "POST", path, map[string]string{"question": "side", "client_request_id": "live"}, 202))
|
||||
snap, err := f.m.pg.SideSnapshot(t.Context(), p.Key())
|
||||
if err != nil || snap == nil || snap.Request.Messages[0].Text() != "live checkpoint" {
|
||||
t.Fatalf("admission checkpoint not durable %+v %v", snap, err)
|
||||
}
|
||||
stopMain()
|
||||
<-done
|
||||
f.call(t, "POST", "/api/side-questions/"+e.ID+"/cancel", nil, 200)
|
||||
waitSide(t, f.m.pg, e.ID, "cancelled")
|
||||
// Simulate an abrupt process exit: persisted running row survives, memory doesn't.
|
||||
orphan, _, err := f.m.pg.StartSideRequest(t.Context(), *snap, "orphan", "restart")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
orphan.Answer = "saved before crash"
|
||||
orphan.Sequence = 1
|
||||
if _, err = f.m.pg.UpdateSideRequest(t.Context(), *orphan); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ctx2, cancel := context.WithCancel(t.Context())
|
||||
defer cancel()
|
||||
restarted := &Server{m: f.m, ctx: ctx2, llmCfg: f.s.llmCfg, llmDirect: bound, llmOn: true}
|
||||
restarted.initSideQuestions()
|
||||
if got := waitSide(t, f.m.pg, orphan.ID, "interrupted"); got.Answer != "saved before crash" {
|
||||
t.Fatal("restart discarded partial answer")
|
||||
}
|
||||
restored, err := restarted.sideSnapshot(t.Context(), p)
|
||||
if err != nil || restored == nil {
|
||||
t.Fatal("restart missing snapshot")
|
||||
}
|
||||
if _, err = restarted.sideProvider(restored.Model); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if req, err := sidequestion.BuildRequest(*restored, nil, "continue after restart"); err != nil || req.Messages[0].Text() != "live checkpoint" {
|
||||
t.Fatal("restart cannot ask from stored context")
|
||||
}
|
||||
// All routes pass the same authentication middleware.
|
||||
w := httptest.NewRecorder()
|
||||
f.handler.ServeHTTP(w, httptest.NewRequest("GET", path, nil))
|
||||
if w.Code != 401 {
|
||||
t.Fatalf("unauthenticated side route: %d", w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSideRejectsDeletedOrChangedCachedProfile(t *testing.T) {
|
||||
f := newSideHTTPFixture(t)
|
||||
p := &db.LLMProfile{Name: "btw-cached-profile", Format: "openai", BaseURL: "http://fixture.invalid", APIKey: "test-only", Model: "fixture", Streaming: true}
|
||||
id, err := f.m.pg.SaveProfile(p)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = f.m.pg.DeleteProfile(id) })
|
||||
cfg, ok := f.s.loadProfileConfig(id)
|
||||
if !ok {
|
||||
t.Fatal("profile unavailable")
|
||||
}
|
||||
f.s.provByProfile = map[int64]*provEntry{id: {prov: f.provider, cfg: cfg}}
|
||||
model := sideModel(cfg, id, p.Name)
|
||||
if _, err = f.s.sideProvider(model); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = f.m.pg.Exec(`UPDATE llm_profiles SET model='changed' WHERE id=$1`, id); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = f.s.sideProvider(model); err == nil {
|
||||
t.Fatal("cached profile bypassed model identity check")
|
||||
}
|
||||
if err = f.m.pg.DeleteProfile(id); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = f.s.sideProvider(model); err == nil {
|
||||
t.Fatal("deleted profile remained usable from cache")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSideTaskDrainPersistsBeforeArchive(t *testing.T) {
|
||||
f := newSideHTTPFixture(t)
|
||||
task, err := f.m.CreateTask("side archive drain", "context", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
id, _ := strconv.ParseInt(task.ID, 10, 64)
|
||||
t.Cleanup(func() { _ = f.m.pg.DeleteTask(id) })
|
||||
p := sidequestion.Parent{TaskID: id, ExplorationID: task.ExpID}
|
||||
f.checkpoint(t, p)
|
||||
path := "/api/tasks/" + task.ID + "/chat/side-questions"
|
||||
e := decodeSide(t, f.call(t, "POST", path, map[string]string{"question": "answer before archive", "client_request_id": "drain"}, 202))
|
||||
select {
|
||||
case <-f.provider.started:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("side never started")
|
||||
}
|
||||
// The real archive entry point closes this same admission barrier first.
|
||||
if !f.s.beginTaskDelete(task.ID) {
|
||||
t.Fatal("cannot close task admission")
|
||||
}
|
||||
defer f.s.abortTaskDelete(task.ID)
|
||||
f.call(t, "POST", path, map[string]string{"question": "too late", "client_request_id": "late"}, 409)
|
||||
ctx, cancel := context.WithTimeout(t.Context(), 3*time.Second)
|
||||
defer cancel()
|
||||
if err := f.s.drainTaskSideQuestions(ctx, task.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
row, err := f.m.pg.SideRequest(t.Context(), e.ID)
|
||||
if err != nil || row == nil || row.Status != "cancelled" || row.Answer != "partial answer" || row.Usage.InputTokens != 19 {
|
||||
t.Fatalf("drain returned before final persistence: %+v %v", row, err)
|
||||
}
|
||||
archive, err := f.m.pg.SnapshotTaskArchive(id)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var rows []sidequestion.Exchange
|
||||
if err := json.Unmarshal(archive.Tables["side_question_requests"], &rows); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(rows) != 1 || rows[0].Status != "cancelled" || rows[0].Usage.InputTokens != 19 {
|
||||
t.Fatalf("archive lost drained side answer/usage: %+v", rows)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSideRestoredWorkerRuntimePublishesNewCheckpoint(t *testing.T) {
|
||||
f := newSideHTTPFixture(t)
|
||||
task, err := f.m.CreateTask("restore side publisher", "context", nil, 3600, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
id, _ := strconv.ParseInt(task.ID, 10, 64)
|
||||
t.Cleanup(func() { _ = f.m.pg.DeleteTask(id) })
|
||||
if err := f.m.pg.SetPaused(id, true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
iid, err := task.Store.AddNode(db.KindIntent, map[string]any{"summary": "restored worker"}, 1, "paused", "planner", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Exercise the production startup path, including the deadline coordinator
|
||||
// which registers the shared runtime before the worker loops.
|
||||
f.s.restoreTaskRuntimes()
|
||||
t.Cleanup(func() {
|
||||
for _, task := range f.m.List() {
|
||||
f.s.engine.StopTask(task.ID)
|
||||
}
|
||||
})
|
||||
f.s.engine.runtimeMu.Lock()
|
||||
runtime := f.s.engine.runtimes[task.ID]
|
||||
f.s.engine.runtimeMu.Unlock()
|
||||
if runtime == nil {
|
||||
t.Fatal("restored task runtime missing")
|
||||
}
|
||||
p := sidequestion.Parent{TaskID: id, ExplorationID: task.ExpID, IntentID: iid}
|
||||
close(f.provider.release)
|
||||
bound := bindSideProvider(f.provider, f.s.llmCfg, 0, "fixture")
|
||||
ctx, deps := sidequestion.Attach(runtime.ctx, p, harness.QueryDeps{}, bound)
|
||||
_, _, _, err = deps.CallModelSync(ctx, llm.CompletionRequest{Messages: []llm.Message{llm.UserText("after restart context")}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
f.s.flushSideSnapshots()
|
||||
snap, err := f.m.pg.SideSnapshot(t.Context(), p.Key())
|
||||
if err != nil || snap == nil || snap.Request.Messages[0].Text() != "after restart context" {
|
||||
t.Fatalf("restored runtime lost checkpoint publisher: %+v %v", snap, err)
|
||||
}
|
||||
}
|
||||
@@ -131,6 +131,9 @@ func (s *Server) archiveTask(job *pgdb.TaskArchive) (runErr error) {
|
||||
if err := s.waitTaskQuiescent(drainCtx, taskID); err != nil {
|
||||
return errors.New("任务仍有运行中的 Agent，请先暂停后重试归档")
|
||||
}
|
||||
if err := s.drainTaskSideQuestions(drainCtx, taskID); err != nil {
|
||||
return err
|
||||
}
|
||||
task, err := s.m.pg.GetTask(job.TaskID)
|
||||
if err != nil {
|
||||
return err
|
||||
|
||||
@@ -204,6 +204,7 @@ func (s *Server) applyIntentControl(ctx context.Context, t *Task, iid int64, act
|
||||
if _, err := t.Store.StopIntentWithReason(iid, reason, "user"); err != nil {
|
||||
return out, err
|
||||
}
|
||||
s.cancelWorkerSide(t.ID, t.ExpID, iid)
|
||||
t.NotifyCancelled(iid, reason)
|
||||
out.State = "stopped"
|
||||
default:
|
||||
|
||||
Reference in new issue
Block a user