mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-13 00:14:33 +00:00
100 lines
3.9 KiB
Go
100 lines
3.9 KiB
Go
package cmd
|
|
|
|
import (
|
|
"context"
|
|
"log/slog"
|
|
|
|
"github.com/google/uuid"
|
|
|
|
"github.com/nextlevelbuilder/goclaw/internal/agent"
|
|
"github.com/nextlevelbuilder/goclaw/internal/bus"
|
|
"github.com/nextlevelbuilder/goclaw/internal/channels"
|
|
"github.com/nextlevelbuilder/goclaw/internal/config"
|
|
"github.com/nextlevelbuilder/goclaw/internal/gateway"
|
|
"github.com/nextlevelbuilder/goclaw/internal/gateway/methods"
|
|
"github.com/nextlevelbuilder/goclaw/internal/heartbeat"
|
|
"github.com/nextlevelbuilder/goclaw/internal/providers"
|
|
"github.com/nextlevelbuilder/goclaw/internal/scheduler"
|
|
"github.com/nextlevelbuilder/goclaw/internal/sessions"
|
|
"github.com/nextlevelbuilder/goclaw/internal/store"
|
|
"github.com/nextlevelbuilder/goclaw/internal/tools"
|
|
"github.com/nextlevelbuilder/goclaw/pkg/protocol"
|
|
)
|
|
|
|
// makeHeartbeatRunFn creates a function that routes a heartbeat run through the scheduler's cron lane.
|
|
func makeHeartbeatRunFn(sched *scheduler.Scheduler) func(ctx context.Context, req agent.RunRequest) <-chan scheduler.RunOutcome {
|
|
return func(ctx context.Context, req agent.RunRequest) <-chan scheduler.RunOutcome {
|
|
return sched.Schedule(ctx, scheduler.LaneCron, req)
|
|
}
|
|
}
|
|
|
|
// startCronAndHeartbeat starts the cron service and heartbeat ticker, wires the heartbeat
|
|
// wake function to the tool + RPC methods, and sets the adaptive token estimate function.
|
|
// Returns the heartbeat ticker (needed by lifecycle for shutdown).
|
|
func startCronAndHeartbeat(
|
|
pgStores *store.Stores,
|
|
server *gateway.Server,
|
|
sched *scheduler.Scheduler,
|
|
msgBus *bus.MessageBus,
|
|
providerRegistry *providers.Registry,
|
|
channelMgr *channels.Manager,
|
|
cfg *config.Config,
|
|
heartbeatTool *tools.HeartbeatTool,
|
|
heartbeatMethods *methods.HeartbeatMethods,
|
|
) *heartbeat.Ticker {
|
|
// Start cron service with job handler (routes through scheduler's cron lane)
|
|
pgStores.Cron.SetOnJob(makeCronJobHandler(sched, msgBus, cfg, channelMgr, pgStores.Sessions, pgStores.Agents, pgStores.Tenants, pgStores.Providers, providerRegistry))
|
|
pgStores.Cron.SetOnEvent(func(event store.CronEvent) {
|
|
server.BroadcastEvent(*protocol.NewEvent(protocol.EventCron, event))
|
|
})
|
|
if err := pgStores.Cron.Start(); err != nil {
|
|
slog.Warn("cron service failed to start", "error", err)
|
|
}
|
|
|
|
// Start heartbeat ticker (routes through scheduler's cron lane)
|
|
heartbeatTicker := heartbeat.NewTicker(heartbeat.TickerConfig{
|
|
Store: pgStores.Heartbeats,
|
|
Agents: pgStores.Agents,
|
|
Sessions: pgStores.Sessions,
|
|
ProviderStore: pgStores.Providers,
|
|
ProviderReg: providerRegistry,
|
|
MsgBus: msgBus,
|
|
Sched: sched,
|
|
RunAgent: makeHeartbeatRunFn(sched),
|
|
ResolveGroupContext: func(ctx context.Context, channel, chatID string) (string, string) {
|
|
channelType := resolveChannelType(channelMgr, channel)
|
|
return channelType, resolveGroupDisplayTitle(ctx, channelMgr, channel, chatID, string(sessions.PeerGroup), "")
|
|
},
|
|
})
|
|
heartbeatTicker.SetOnEvent(func(event store.HeartbeatEvent) {
|
|
server.BroadcastEvent(*protocol.NewEvent(protocol.EventHeartbeat, event))
|
|
})
|
|
heartbeatTicker.Start()
|
|
|
|
// Wire heartbeat wake function to tool + RPC + cron wakeMode
|
|
heartbeatTool.SetWakeFn(heartbeatTicker.Wake)
|
|
heartbeatMethods.SetWakeFn(heartbeatTicker.Wake)
|
|
heartbeatMethods.SetAgentStore(pgStores.Agents)
|
|
heartbeatMethods.SetProviderStore(pgStores.Providers)
|
|
cronHeartbeatWakeFn = func(agentID string) {
|
|
if id, err := uuid.Parse(agentID); err == nil {
|
|
heartbeatTicker.Wake(id)
|
|
}
|
|
}
|
|
|
|
// Adaptive throttle: reduce per-session concurrency when nearing the summary threshold.
|
|
sched.SetTokenEstimateFunc(func(sessionKey string) (int, int) {
|
|
bctx := context.Background()
|
|
history := pgStores.Sessions.GetHistory(bctx, sessionKey)
|
|
lastPT, lastMC := pgStores.Sessions.GetLastPromptTokens(bctx, sessionKey)
|
|
tokens := agent.EstimateTokensWithCalibration(history, lastPT, lastMC)
|
|
cw := pgStores.Sessions.GetContextWindow(bctx, sessionKey)
|
|
if cw <= 0 {
|
|
cw = config.DefaultContextWindow
|
|
}
|
|
return tokens, cw
|
|
})
|
|
|
|
return heartbeatTicker
|
|
}
|