Merge pull request #43 from dj-oyu/upstream-sync-20260315
upstream sync + hook consolidation refactor
This commit is contained in:
commit
46c0d1443d
11 changed files with 2737 additions and 1404 deletions
|
|
@ -89,32 +89,14 @@ func (cb *ContextBuilder) getIdentity() string {
|
||||||
// Build tools section dynamically
|
// Build tools section dynamically
|
||||||
toolsSection := cb.buildToolsSection()
|
toolsSection := cb.buildToolsSection()
|
||||||
|
|
||||||
// Build prompt with optional orchestration banner
|
// Orchestration overrides (banner, identity, executing rule)
|
||||||
var prompt string
|
prompt, identity, executingRule := cb.extIdentityOverrides()
|
||||||
if cb.orchestrationEnabled {
|
if identity == "" {
|
||||||
prompt = ` /_/_/_/_/_/_/_/_/_/_/_/_/_/_/
|
identity = "a helpful AI assistant"
|
||||||
|
|
||||||
O R C H E S T R A M O D E
|
|
||||||
|
|
||||||
/_/_/_/_/_/_/_/_/_/_/_/_/_/_/
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
`
|
|
||||||
}
|
}
|
||||||
|
if executingRule == "" {
|
||||||
// Conditional identity and plan executing rule for orchestration mode
|
executingRule = `Work through the current Phase's steps.
|
||||||
identity := "a helpful AI assistant"
|
|
||||||
executingRule := `Work through the current Phase's steps.
|
|
||||||
Mark each "- [x]" via edit_file. The system will auto-advance phases.`
|
Mark each "- [x]" via edit_file. The system will auto-advance phases.`
|
||||||
if cb.orchestrationEnabled {
|
|
||||||
identity = "a conductor AI agent that orchestrates subagents"
|
|
||||||
executingRule = `Delegate the current Phase's steps to subagents using spawn.
|
|
||||||
For each step: spawn a subagent with the appropriate preset (scout for investigation,
|
|
||||||
coder for implementation, analyst for review). Spawn multiple independent steps in parallel.
|
|
||||||
When a subagent completes, mark "- [x]" via edit_file and record findings in
|
|
||||||
## Orchestration > Findings in MEMORY.md.
|
|
||||||
Only do a step inline if it's a single quick tool call (e.g., reading one file).`
|
|
||||||
}
|
}
|
||||||
|
|
||||||
return fmt.Sprintf(prompt+`# picoclaw 🦞 (%s)
|
return fmt.Sprintf(prompt+`# picoclaw 🦞 (%s)
|
||||||
|
|
@ -175,29 +157,6 @@ Your workspace is at: %s
|
||||||
toolsSection, executingRule, toolDiscovery)
|
toolsSection, executingRule, toolDiscovery)
|
||||||
}
|
}
|
||||||
|
|
||||||
func (cb *ContextBuilder) buildToolsSection() string {
|
|
||||||
if cb.tools == nil {
|
|
||||||
return ""
|
|
||||||
}
|
|
||||||
|
|
||||||
summaries := cb.tools.GetSummaries()
|
|
||||||
if len(summaries) == 0 {
|
|
||||||
return ""
|
|
||||||
}
|
|
||||||
|
|
||||||
var sb strings.Builder
|
|
||||||
sb.WriteString("## Available Tools\n\n")
|
|
||||||
sb.WriteString(
|
|
||||||
"**CRITICAL**: You MUST use tools to perform actions. Do NOT pretend to execute commands or schedule tasks.\n\n",
|
|
||||||
)
|
|
||||||
sb.WriteString("You have access to the following tools:\n\n")
|
|
||||||
for _, s := range summaries {
|
|
||||||
sb.WriteString(s)
|
|
||||||
sb.WriteString("\n")
|
|
||||||
}
|
|
||||||
return sb.String()
|
|
||||||
}
|
|
||||||
|
|
||||||
func (cb *ContextBuilder) getDiscoveryRule() string {
|
func (cb *ContextBuilder) getDiscoveryRule() string {
|
||||||
if !cb.toolDiscoveryBM25 && !cb.toolDiscoveryRegex {
|
if !cb.toolDiscoveryBM25 && !cb.toolDiscoveryRegex {
|
||||||
return ""
|
return ""
|
||||||
|
|
@ -223,12 +182,8 @@ func (cb *ContextBuilder) BuildSystemPrompt() string {
|
||||||
// Core identity section
|
// Core identity section
|
||||||
parts = append(parts, cb.getIdentity())
|
parts = append(parts, cb.getIdentity())
|
||||||
|
|
||||||
// Orchestration guidance — injected only when spawn tool is registered
|
// Fork-specific prompt sections (orchestration guidance, peer note)
|
||||||
if cb.tools != nil {
|
parts = append(parts, cb.extPromptSections()...)
|
||||||
if _, hasSpawn := cb.tools.Get("spawn"); hasSpawn {
|
|
||||||
parts = append(parts, orchestrationGuidance)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Bootstrap files
|
// Bootstrap files
|
||||||
bootstrapContent := cb.LoadBootstrapFiles()
|
bootstrapContent := cb.LoadBootstrapFiles()
|
||||||
|
|
@ -253,11 +208,6 @@ The following skills extend your capabilities. To use a skill, read its SKILL.md
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Peer session coordination
|
|
||||||
if cb.peerNote != "" {
|
|
||||||
parts = append(parts, "## Active Sessions\n\n"+cb.peerNote)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Memory context
|
// Memory context
|
||||||
memoryContext := cb.memory.GetMemoryContext()
|
memoryContext := cb.memory.GetMemoryContext()
|
||||||
if memoryContext != "" {
|
if memoryContext != "" {
|
||||||
|
|
|
||||||
|
|
@ -1,6 +1,10 @@
|
||||||
package agent
|
package agent
|
||||||
|
|
||||||
import "github.com/sipeed/picoclaw/pkg/tools"
|
import (
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/tools"
|
||||||
|
)
|
||||||
|
|
||||||
// contextBuilderExt holds fork-specific fields for ContextBuilder.
|
// contextBuilderExt holds fork-specific fields for ContextBuilder.
|
||||||
// Embedded in ContextBuilder so existing field access (cb.workDir, cb.tools, etc.) continues to work.
|
// Embedded in ContextBuilder so existing field access (cb.workDir, cb.tools, etc.) continues to work.
|
||||||
|
|
@ -33,6 +37,77 @@ func (cb *ContextBuilder) SetOrchestrationEnabled(enabled bool) {
|
||||||
cb.orchestrationEnabled = enabled
|
cb.orchestrationEnabled = enabled
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// buildToolsSection generates the "Available Tools" section for the system prompt.
|
||||||
|
func (cb *ContextBuilder) buildToolsSection() string {
|
||||||
|
if cb.tools == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
|
||||||
|
summaries := cb.tools.GetSummaries()
|
||||||
|
if len(summaries) == 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
|
||||||
|
var sb strings.Builder
|
||||||
|
sb.WriteString("## Available Tools\n\n")
|
||||||
|
sb.WriteString(
|
||||||
|
"**CRITICAL**: You MUST use tools to perform actions. Do NOT pretend to execute commands or schedule tasks.\n\n",
|
||||||
|
)
|
||||||
|
sb.WriteString("You have access to the following tools:\n\n")
|
||||||
|
for _, s := range summaries {
|
||||||
|
sb.WriteString(s)
|
||||||
|
sb.WriteString("\n")
|
||||||
|
}
|
||||||
|
return sb.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
// extIdentityOverrides returns the orchestration-specific overrides for
|
||||||
|
// getIdentity: banner prefix, identity string, and plan executing rule.
|
||||||
|
// When orchestration is disabled, all return values are empty strings.
|
||||||
|
func (cb *ContextBuilder) extIdentityOverrides() (banner, identity, executingRule string) {
|
||||||
|
if !cb.orchestrationEnabled {
|
||||||
|
return "", "", ""
|
||||||
|
}
|
||||||
|
|
||||||
|
banner = ` /_/_/_/_/_/_/_/_/_/_/_/_/_/_/
|
||||||
|
|
||||||
|
O R C H E S T R A M O D E
|
||||||
|
|
||||||
|
/_/_/_/_/_/_/_/_/_/_/_/_/_/_/
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
`
|
||||||
|
identity = "a conductor AI agent that orchestrates subagents"
|
||||||
|
executingRule = `Delegate the current Phase's steps to subagents using spawn.
|
||||||
|
For each step: spawn a subagent with the appropriate preset (scout for investigation,
|
||||||
|
coder for implementation, analyst for review). Spawn multiple independent steps in parallel.
|
||||||
|
When a subagent completes, mark "- [x]" via edit_file and record findings in
|
||||||
|
## Orchestration > Findings in MEMORY.md.
|
||||||
|
Only do a step inline if it's a single quick tool call (e.g., reading one file).`
|
||||||
|
return banner, identity, executingRule
|
||||||
|
}
|
||||||
|
|
||||||
|
// extPromptSections returns fork-specific prompt sections to append to
|
||||||
|
// BuildSystemPrompt: orchestration guidance and peer session note.
|
||||||
|
func (cb *ContextBuilder) extPromptSections() []string {
|
||||||
|
var sections []string
|
||||||
|
|
||||||
|
// Orchestration guidance — injected only when spawn tool is registered
|
||||||
|
if cb.tools != nil {
|
||||||
|
if _, hasSpawn := cb.tools.Get("spawn"); hasSpawn {
|
||||||
|
sections = append(sections, orchestrationGuidance)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Peer session coordination
|
||||||
|
if cb.peerNote != "" {
|
||||||
|
sections = append(sections, "## Active Sessions\n\n"+cb.peerNote)
|
||||||
|
}
|
||||||
|
|
||||||
|
return sections
|
||||||
|
}
|
||||||
|
|
||||||
// Memory returns the underlying MemoryStore for direct plan queries.
|
// Memory returns the underlying MemoryStore for direct plan queries.
|
||||||
func (cb *ContextBuilder) Memory() *MemoryStore {
|
func (cb *ContextBuilder) Memory() *MemoryStore {
|
||||||
return cb.memory
|
return cb.memory
|
||||||
|
|
|
||||||
|
|
@ -9,7 +9,6 @@ import (
|
||||||
"strings"
|
"strings"
|
||||||
|
|
||||||
"github.com/sipeed/picoclaw/pkg/config"
|
"github.com/sipeed/picoclaw/pkg/config"
|
||||||
"github.com/sipeed/picoclaw/pkg/git"
|
|
||||||
"github.com/sipeed/picoclaw/pkg/providers"
|
"github.com/sipeed/picoclaw/pkg/providers"
|
||||||
"github.com/sipeed/picoclaw/pkg/routing"
|
"github.com/sipeed/picoclaw/pkg/routing"
|
||||||
"github.com/sipeed/picoclaw/pkg/session"
|
"github.com/sipeed/picoclaw/pkg/session"
|
||||||
|
|
@ -38,6 +37,8 @@ type AgentInstance struct {
|
||||||
Sessions *session.LegacyAdapter
|
Sessions *session.LegacyAdapter
|
||||||
ContextBuilder *ContextBuilder
|
ContextBuilder *ContextBuilder
|
||||||
Tools *tools.ToolRegistry
|
Tools *tools.ToolRegistry
|
||||||
|
Subagents *config.SubagentsConfig
|
||||||
|
SkillsFilter []string
|
||||||
Candidates []providers.FallbackCandidate
|
Candidates []providers.FallbackCandidate
|
||||||
PlanModel string
|
PlanModel string
|
||||||
PlanFallbacks []string
|
PlanFallbacks []string
|
||||||
|
|
@ -141,23 +142,10 @@ func NewAgentInstance(
|
||||||
|
|
||||||
agentID := routing.DefaultAgentID
|
agentID := routing.DefaultAgentID
|
||||||
agentName := ""
|
agentName := ""
|
||||||
var subagents *config.SubagentsConfig
|
|
||||||
var skillsFilter []string
|
|
||||||
|
|
||||||
if agentCfg != nil {
|
if agentCfg != nil {
|
||||||
agentID = routing.NormalizeAgentID(agentCfg.ID)
|
agentID = routing.NormalizeAgentID(agentCfg.ID)
|
||||||
agentName = agentCfg.Name
|
agentName = agentCfg.Name
|
||||||
subagents = agentCfg.Subagents
|
|
||||||
skillsFilter = agentCfg.Skills
|
|
||||||
}
|
|
||||||
|
|
||||||
// Apply defaults.Orchestration: if the flag is set, ensure orchestration is enabled.
|
|
||||||
if defaults.Orchestration {
|
|
||||||
if subagents == nil {
|
|
||||||
subagents = &config.SubagentsConfig{Enabled: true}
|
|
||||||
} else {
|
|
||||||
subagents.Enabled = true
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
maxIter := defaults.MaxToolIterations
|
maxIter := defaults.MaxToolIterations
|
||||||
|
|
@ -243,19 +231,6 @@ func NewAgentInstance(
|
||||||
|
|
||||||
candidates := providers.ResolveCandidatesWithLookup(modelCfg, defaults.Provider, resolveFromModelList)
|
candidates := providers.ResolveCandidatesWithLookup(modelCfg, defaults.Provider, resolveFromModelList)
|
||||||
|
|
||||||
// Resolve plan model (for interviewing/review phases)
|
|
||||||
planModel := resolvePlanModel(agentCfg, defaults)
|
|
||||||
planFallbacks := resolvePlanFallbacks(agentCfg, defaults)
|
|
||||||
|
|
||||||
var planCandidates []providers.FallbackCandidate
|
|
||||||
if planModel != "" {
|
|
||||||
planModelCfg := providers.ModelConfig{
|
|
||||||
Primary: planModel,
|
|
||||||
Fallbacks: planFallbacks,
|
|
||||||
}
|
|
||||||
planCandidates = providers.ResolveCandidates(planModelCfg, defaults.Provider)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Model routing setup: pre-resolve light model candidates at creation time
|
// Model routing setup: pre-resolve light model candidates at creation time
|
||||||
// to avoid repeated model_list lookups on every incoming message.
|
// to avoid repeated model_list lookups on every incoming message.
|
||||||
var router *routing.Router
|
var router *routing.Router
|
||||||
|
|
@ -275,17 +250,7 @@ func NewAgentInstance(
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Startup cleanup: prune orphaned worktrees
|
agent := &AgentInstance{
|
||||||
worktreesDir := filepath.Join(workspace, ".worktrees")
|
|
||||||
if repoRoot := git.FindRepoRoot(workspace); repoRoot != "" {
|
|
||||||
git.PruneOrphaned(repoRoot, worktreesDir)
|
|
||||||
}
|
|
||||||
|
|
||||||
return &AgentInstance{
|
|
||||||
instanceExt: instanceExt{
|
|
||||||
Subagents: subagents,
|
|
||||||
SkillsFilter: skillsFilter,
|
|
||||||
},
|
|
||||||
ID: agentID,
|
ID: agentID,
|
||||||
Name: agentName,
|
Name: agentName,
|
||||||
Model: model,
|
Model: model,
|
||||||
|
|
@ -304,12 +269,14 @@ func NewAgentInstance(
|
||||||
ContextBuilder: contextBuilder,
|
ContextBuilder: contextBuilder,
|
||||||
Tools: toolsRegistry,
|
Tools: toolsRegistry,
|
||||||
Candidates: candidates,
|
Candidates: candidates,
|
||||||
PlanModel: planModel,
|
|
||||||
PlanFallbacks: planFallbacks,
|
|
||||||
PlanCandidates: planCandidates,
|
|
||||||
Router: router,
|
Router: router,
|
||||||
LightCandidates: lightCandidates,
|
LightCandidates: lightCandidates,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Initialize fork-specific fields (subagents, plan model, worktree pruning).
|
||||||
|
agent.initInstanceExt(agentCfg, defaults, cfg)
|
||||||
|
|
||||||
|
return agent
|
||||||
}
|
}
|
||||||
|
|
||||||
// resolveAgentWorkspace determines the workspace directory for an agent.
|
// resolveAgentWorkspace determines the workspace directory for an agent.
|
||||||
|
|
|
||||||
|
|
@ -2,10 +2,12 @@ package agent
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"path/filepath"
|
||||||
"sync"
|
"sync"
|
||||||
|
|
||||||
"github.com/sipeed/picoclaw/pkg/config"
|
"github.com/sipeed/picoclaw/pkg/config"
|
||||||
"github.com/sipeed/picoclaw/pkg/git"
|
"github.com/sipeed/picoclaw/pkg/git"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/providers"
|
||||||
"github.com/sipeed/picoclaw/pkg/tools"
|
"github.com/sipeed/picoclaw/pkg/tools"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -16,9 +18,6 @@ type instanceExt struct {
|
||||||
// Used by runAgentLoop to wait for spawned subagents before worktree cleanup.
|
// Used by runAgentLoop to wait for spawned subagents before worktree cleanup.
|
||||||
SubagentMgr *tools.SubagentManager
|
SubagentMgr *tools.SubagentManager
|
||||||
|
|
||||||
Subagents *config.SubagentsConfig
|
|
||||||
SkillsFilter []string
|
|
||||||
|
|
||||||
// Interview staleness tracking: consecutive turns where MEMORY.md was not updated.
|
// Interview staleness tracking: consecutive turns where MEMORY.md was not updated.
|
||||||
interviewStaleCount int
|
interviewStaleCount int
|
||||||
interviewMemoryLen int
|
interviewMemoryLen int
|
||||||
|
|
@ -28,6 +27,47 @@ type instanceExt struct {
|
||||||
worktreeMu sync.RWMutex
|
worktreeMu sync.RWMutex
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// initInstanceExt initializes fork-specific fields: subagents config,
|
||||||
|
// skills filter, plan model resolution, and worktree pruning.
|
||||||
|
func (ai *AgentInstance) initInstanceExt(
|
||||||
|
agentCfg *config.AgentConfig,
|
||||||
|
defaults *config.AgentDefaults,
|
||||||
|
cfg *config.Config,
|
||||||
|
) {
|
||||||
|
// Extract subagents and skills filter from agent config
|
||||||
|
if agentCfg != nil {
|
||||||
|
ai.Subagents = agentCfg.Subagents
|
||||||
|
ai.SkillsFilter = agentCfg.Skills
|
||||||
|
}
|
||||||
|
|
||||||
|
// Apply defaults.Orchestration: if the flag is set, ensure orchestration is enabled.
|
||||||
|
if defaults.Orchestration {
|
||||||
|
if ai.Subagents == nil {
|
||||||
|
ai.Subagents = &config.SubagentsConfig{Enabled: true}
|
||||||
|
} else {
|
||||||
|
ai.Subagents.Enabled = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Resolve plan model (for interviewing/review phases)
|
||||||
|
ai.PlanModel = resolvePlanModel(agentCfg, defaults)
|
||||||
|
ai.PlanFallbacks = resolvePlanFallbacks(agentCfg, defaults)
|
||||||
|
|
||||||
|
if ai.PlanModel != "" {
|
||||||
|
planModelCfg := providers.ModelConfig{
|
||||||
|
Primary: ai.PlanModel,
|
||||||
|
Fallbacks: ai.PlanFallbacks,
|
||||||
|
}
|
||||||
|
ai.PlanCandidates = providers.ResolveCandidates(planModelCfg, defaults.Provider)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Startup cleanup: prune orphaned worktrees
|
||||||
|
worktreesDir := filepath.Join(ai.Workspace, ".worktrees")
|
||||||
|
if repoRoot := git.FindRepoRoot(ai.Workspace); repoRoot != "" {
|
||||||
|
git.PruneOrphaned(repoRoot, worktreesDir)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// ActivateWorktree creates a worktree for a session.
|
// ActivateWorktree creates a worktree for a session.
|
||||||
// projectDir is the git repository to create the worktree in.
|
// projectDir is the git repository to create the worktree in.
|
||||||
// If empty, falls back to ai.Workspace.
|
// If empty, falls back to ai.Workspace.
|
||||||
|
|
|
||||||
1312
pkg/agent/loop.go
1312
pkg/agent/loop.go
File diff suppressed because it is too large
Load diff
|
|
@ -1,11 +1,18 @@
|
||||||
package agent
|
package agent
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"strings"
|
||||||
"sync"
|
"sync"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/bus"
|
||||||
"github.com/sipeed/picoclaw/pkg/config"
|
"github.com/sipeed/picoclaw/pkg/config"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/logger"
|
||||||
"github.com/sipeed/picoclaw/pkg/orch"
|
"github.com/sipeed/picoclaw/pkg/orch"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/providers"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/routing"
|
||||||
"github.com/sipeed/picoclaw/pkg/stats"
|
"github.com/sipeed/picoclaw/pkg/stats"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/tools"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/utils"
|
||||||
)
|
)
|
||||||
|
|
||||||
// loopExt holds fork-specific fields for AgentLoop.
|
// loopExt holds fork-specific fields for AgentLoop.
|
||||||
|
|
@ -35,6 +42,61 @@ type loopExt struct {
|
||||||
onHeartbeatThreadUpdate func(int)
|
onHeartbeatThreadUpdate func(int)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// initLoopExt initializes all fork-specific fields: stats tracker,
|
||||||
|
// session tracker, orchestration broadcaster, and background goroutines.
|
||||||
|
// Called from NewAgentLoop after the struct is constructed.
|
||||||
|
func (al *AgentLoop) initLoopExt(cfg *config.Config, registry *AgentRegistry, enableStats bool) {
|
||||||
|
defaultAgent := registry.GetDefaultAgent()
|
||||||
|
|
||||||
|
// Stats tracker
|
||||||
|
if enableStats && defaultAgent != nil {
|
||||||
|
al.stats = stats.NewTracker(defaultAgent.Workspace)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Session tracker
|
||||||
|
al.sessions = NewSessionTracker()
|
||||||
|
|
||||||
|
// Orchestration broadcaster — needed if any agent has subagents enabled.
|
||||||
|
// Note: instance.go maps defaults.Orchestration → Subagents.Enabled,
|
||||||
|
// so --orchestration is automatically reflected here.
|
||||||
|
al.orchReporter = orch.Noop
|
||||||
|
for _, id := range registry.ListAgentIDs() {
|
||||||
|
if a, ok := registry.GetAgent(id); ok && a.Subagents != nil && a.Subagents.Enabled {
|
||||||
|
al.orchBroadcaster = orch.NewBroadcaster()
|
||||||
|
al.orchReporter = al.orchBroadcaster
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Shutdown signal channel
|
||||||
|
al.done = make(chan struct{})
|
||||||
|
|
||||||
|
// Background GC goroutine
|
||||||
|
go al.gcLoop()
|
||||||
|
}
|
||||||
|
|
||||||
|
// closeExt releases fork-specific resources: done channel, stats tracker,
|
||||||
|
// and session stores for all agents.
|
||||||
|
func (al *AgentLoop) closeExt() {
|
||||||
|
select {
|
||||||
|
case <-al.done:
|
||||||
|
// already closed
|
||||||
|
default:
|
||||||
|
close(al.done)
|
||||||
|
}
|
||||||
|
|
||||||
|
if al.stats != nil {
|
||||||
|
al.stats.Close()
|
||||||
|
}
|
||||||
|
|
||||||
|
registry := al.GetRegistry()
|
||||||
|
for _, agentID := range registry.ListAgentIDs() {
|
||||||
|
if agent, ok := registry.GetAgent(agentID); ok {
|
||||||
|
agent.Sessions.Close()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// SetConfigSaver registers a callback to persist config changes.
|
// SetConfigSaver registers a callback to persist config changes.
|
||||||
func (al *AgentLoop) SetConfigSaver(fn func(*config.Config) error) {
|
func (al *AgentLoop) SetConfigSaver(fn func(*config.Config) error) {
|
||||||
al.saveConfig = fn
|
al.saveConfig = fn
|
||||||
|
|
@ -44,3 +106,139 @@ func (al *AgentLoop) SetConfigSaver(fn func(*config.Config) error) {
|
||||||
func (al *AgentLoop) SetHeartbeatThreadUpdater(fn func(int)) {
|
func (al *AgentLoop) SetHeartbeatThreadUpdater(fn func(int)) {
|
||||||
al.onHeartbeatThreadUpdate = fn
|
al.onHeartbeatThreadUpdate = fn
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// registerOrchestrationTools registers spawn, subagent, answer, and review_plan
|
||||||
|
// tools for agents with orchestration enabled.
|
||||||
|
func registerOrchestrationTools(
|
||||||
|
cfg *config.Config,
|
||||||
|
agent *AgentInstance,
|
||||||
|
agentID string,
|
||||||
|
registry *AgentRegistry,
|
||||||
|
provider providers.LLMProvider,
|
||||||
|
msgBus *bus.MessageBus,
|
||||||
|
al *AgentLoop,
|
||||||
|
) {
|
||||||
|
if agent.Subagents == nil || !agent.Subagents.Enabled {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
webSearchOpts := tools.WebSearchToolOptions{
|
||||||
|
BraveAPIKeys: config.MergeAPIKeys(cfg.Tools.Web.Brave.APIKey, cfg.Tools.Web.Brave.APIKeys),
|
||||||
|
BraveMaxResults: cfg.Tools.Web.Brave.MaxResults,
|
||||||
|
BraveEnabled: cfg.Tools.Web.Brave.Enabled,
|
||||||
|
TavilyAPIKeys: config.MergeAPIKeys(cfg.Tools.Web.Tavily.APIKey, cfg.Tools.Web.Tavily.APIKeys),
|
||||||
|
TavilyBaseURL: cfg.Tools.Web.Tavily.BaseURL,
|
||||||
|
TavilyMaxResults: cfg.Tools.Web.Tavily.MaxResults,
|
||||||
|
TavilyEnabled: cfg.Tools.Web.Tavily.Enabled,
|
||||||
|
DuckDuckGoMaxResults: cfg.Tools.Web.DuckDuckGo.MaxResults,
|
||||||
|
DuckDuckGoEnabled: cfg.Tools.Web.DuckDuckGo.Enabled,
|
||||||
|
PerplexityAPIKeys: config.MergeAPIKeys(
|
||||||
|
cfg.Tools.Web.Perplexity.APIKey,
|
||||||
|
cfg.Tools.Web.Perplexity.APIKeys,
|
||||||
|
),
|
||||||
|
PerplexityMaxResults: cfg.Tools.Web.Perplexity.MaxResults,
|
||||||
|
PerplexityEnabled: cfg.Tools.Web.Perplexity.Enabled,
|
||||||
|
}
|
||||||
|
|
||||||
|
subagentManager := tools.NewSubagentManager(
|
||||||
|
provider,
|
||||||
|
agent.Model,
|
||||||
|
agent.Workspace,
|
||||||
|
msgBus,
|
||||||
|
al.reporter(),
|
||||||
|
webSearchOpts,
|
||||||
|
)
|
||||||
|
|
||||||
|
subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature)
|
||||||
|
|
||||||
|
// Wire session recorder for DAG persistence.
|
||||||
|
recorder := newSessionRecorder(agent.Sessions)
|
||||||
|
conductorKey := routing.BuildAgentMainSessionKey(agent.ID)
|
||||||
|
subagentManager.SetSessionRecorder(recorder, conductorKey)
|
||||||
|
|
||||||
|
agent.SubagentMgr = subagentManager
|
||||||
|
|
||||||
|
spawnTool := tools.NewSpawnTool(subagentManager)
|
||||||
|
currentAgentID := agentID
|
||||||
|
spawnTool.SetAllowlistChecker(func(targetAgentID string) bool {
|
||||||
|
return registry.CanSpawnSubagent(currentAgentID, targetAgentID)
|
||||||
|
})
|
||||||
|
agent.Tools.Register(spawnTool)
|
||||||
|
|
||||||
|
// Register blocking subagent tool alongside spawn
|
||||||
|
agent.Tools.Register(tools.NewSubagentTool(subagentManager))
|
||||||
|
|
||||||
|
// Register conductor-side escalation tools (answer questions, review plans)
|
||||||
|
agent.Tools.Register(tools.NewAnswerSubagentTool(subagentManager))
|
||||||
|
agent.Tools.Register(tools.NewReviewSubagentPlanTool(subagentManager))
|
||||||
|
|
||||||
|
// Set orchestration mode on context builder
|
||||||
|
agent.ContextBuilder.SetOrchestrationEnabled(true)
|
||||||
|
}
|
||||||
|
|
||||||
|
// handleTaskIntervention checks if a message is a reply to an active task and
|
||||||
|
// either cancels the task or injects a user intervention. Returns (response, handled).
|
||||||
|
func (al *AgentLoop) handleTaskIntervention(msg bus.InboundMessage) (string, bool) {
|
||||||
|
taskID, ok := msg.Metadata["task_id"]
|
||||||
|
if !ok || taskID == "" {
|
||||||
|
return "", false
|
||||||
|
}
|
||||||
|
|
||||||
|
val, found := al.activeTasks.Load(taskID)
|
||||||
|
if !found {
|
||||||
|
// Task not found — fall through to normal processing
|
||||||
|
return "", false
|
||||||
|
}
|
||||||
|
|
||||||
|
task := val.(*activeTask)
|
||||||
|
|
||||||
|
content := strings.TrimSpace(msg.Content)
|
||||||
|
lower := strings.ToLower(content)
|
||||||
|
|
||||||
|
// Check for stop keywords
|
||||||
|
stopKeywords := []string{
|
||||||
|
"stop", "cancel", "abort",
|
||||||
|
"停止", "中止", "やめて", //nolint:gosmopolitan // intentional CJK stop words
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, kw := range stopKeywords {
|
||||||
|
if lower == kw {
|
||||||
|
task.cancel()
|
||||||
|
|
||||||
|
logger.InfoCF("agent", "Task canceled by user intervention",
|
||||||
|
map[string]any{"task_id": taskID})
|
||||||
|
|
||||||
|
return "Task canceled.", true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Inject message into interrupt channel for the tool loop
|
||||||
|
select {
|
||||||
|
case task.interrupt <- content:
|
||||||
|
logger.InfoCF("agent", "User intervention queued",
|
||||||
|
map[string]any{"task_id": taskID, "content": utils.Truncate(content, 80)})
|
||||||
|
default:
|
||||||
|
logger.WarnCF("agent", "Interrupt channel full, message dropped",
|
||||||
|
map[string]any{"task_id": taskID})
|
||||||
|
}
|
||||||
|
|
||||||
|
return "Intervention sent to running task.", true
|
||||||
|
}
|
||||||
|
|
||||||
|
// expandForkCommands expands fork-specific /skill and /plan commands in the message.
|
||||||
|
// Returns the modified message and the compact form for history.
|
||||||
|
func (al *AgentLoop) expandForkCommands(msg *bus.InboundMessage) string {
|
||||||
|
var expansionCompact string
|
||||||
|
|
||||||
|
if expanded, compact, ok := al.expandSkillCommand(*msg); ok {
|
||||||
|
msg.Content = expanded
|
||||||
|
expansionCompact = compact
|
||||||
|
}
|
||||||
|
|
||||||
|
if expanded, compact, ok := al.expandPlanCommand(*msg); ok {
|
||||||
|
msg.Content = expanded
|
||||||
|
expansionCompact = compact
|
||||||
|
}
|
||||||
|
|
||||||
|
return expansionCompact
|
||||||
|
}
|
||||||
|
|
|
||||||
1073
pkg/agent/loop_run.go
Normal file
1073
pkg/agent/loop_run.go
Normal file
File diff suppressed because it is too large
Load diff
150
pkg/providers/azure/provider.go
Normal file
150
pkg/providers/azure/provider.go
Normal file
|
|
@ -0,0 +1,150 @@
|
||||||
|
package azure
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
"net/url"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/providers/common"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/providers/protocoltypes"
|
||||||
|
)
|
||||||
|
|
||||||
|
type (
|
||||||
|
LLMResponse = protocoltypes.LLMResponse
|
||||||
|
Message = protocoltypes.Message
|
||||||
|
ToolDefinition = protocoltypes.ToolDefinition
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
// azureAPIVersion is the Azure OpenAI API version used for all requests.
|
||||||
|
azureAPIVersion = "2024-10-21"
|
||||||
|
defaultRequestTimeout = common.DefaultRequestTimeout
|
||||||
|
)
|
||||||
|
|
||||||
|
// Provider implements the LLM provider interface for Azure OpenAI endpoints.
|
||||||
|
// It handles Azure-specific authentication (api-key header), URL construction
|
||||||
|
// (deployment-based), and request body formatting (max_completion_tokens, no model field).
|
||||||
|
type Provider struct {
|
||||||
|
apiKey string
|
||||||
|
apiBase string
|
||||||
|
httpClient *http.Client
|
||||||
|
}
|
||||||
|
|
||||||
|
// Option configures the Azure Provider.
|
||||||
|
type Option func(*Provider)
|
||||||
|
|
||||||
|
// WithRequestTimeout sets the HTTP request timeout.
|
||||||
|
func WithRequestTimeout(timeout time.Duration) Option {
|
||||||
|
return func(p *Provider) {
|
||||||
|
if timeout > 0 {
|
||||||
|
p.httpClient.Timeout = timeout
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewProvider creates a new Azure OpenAI provider.
|
||||||
|
func NewProvider(apiKey, apiBase, proxy string, opts ...Option) *Provider {
|
||||||
|
p := &Provider{
|
||||||
|
apiKey: apiKey,
|
||||||
|
apiBase: strings.TrimRight(apiBase, "/"),
|
||||||
|
httpClient: common.NewHTTPClient(proxy),
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, opt := range opts {
|
||||||
|
if opt != nil {
|
||||||
|
opt(p)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return p
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewProviderWithTimeout creates a new Azure OpenAI provider with a custom request timeout in seconds.
|
||||||
|
func NewProviderWithTimeout(apiKey, apiBase, proxy string, requestTimeoutSeconds int) *Provider {
|
||||||
|
return NewProvider(
|
||||||
|
apiKey, apiBase, proxy,
|
||||||
|
WithRequestTimeout(time.Duration(requestTimeoutSeconds)*time.Second),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Chat sends a chat completion request to the Azure OpenAI endpoint.
|
||||||
|
// The model parameter is used as the Azure deployment name in the URL.
|
||||||
|
func (p *Provider) Chat(
|
||||||
|
ctx context.Context,
|
||||||
|
messages []Message,
|
||||||
|
tools []ToolDefinition,
|
||||||
|
model string,
|
||||||
|
options map[string]any,
|
||||||
|
) (*LLMResponse, error) {
|
||||||
|
if p.apiBase == "" {
|
||||||
|
return nil, fmt.Errorf("Azure API base not configured")
|
||||||
|
}
|
||||||
|
|
||||||
|
// model is the deployment name for Azure OpenAI
|
||||||
|
deployment := model
|
||||||
|
|
||||||
|
// Build Azure-specific URL safely using url.JoinPath and query encoding
|
||||||
|
// to prevent path traversal or query injection via deployment names.
|
||||||
|
base, err := url.JoinPath(p.apiBase, "openai/deployments", deployment, "chat/completions")
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to build Azure request URL: %w", err)
|
||||||
|
}
|
||||||
|
requestURL := base + "?api-version=" + azureAPIVersion
|
||||||
|
|
||||||
|
// Build request body — no "model" field (Azure infers from deployment URL)
|
||||||
|
requestBody := map[string]any{
|
||||||
|
"messages": common.SerializeMessages(messages),
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(tools) > 0 {
|
||||||
|
requestBody["tools"] = tools
|
||||||
|
requestBody["tool_choice"] = "auto"
|
||||||
|
}
|
||||||
|
|
||||||
|
// Azure OpenAI always uses max_completion_tokens
|
||||||
|
if maxTokens, ok := common.AsInt(options["max_tokens"]); ok {
|
||||||
|
requestBody["max_completion_tokens"] = maxTokens
|
||||||
|
}
|
||||||
|
|
||||||
|
if temperature, ok := common.AsFloat(options["temperature"]); ok {
|
||||||
|
requestBody["temperature"] = temperature
|
||||||
|
}
|
||||||
|
|
||||||
|
jsonData, err := json.Marshal(requestBody)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to marshal request: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
req, err := http.NewRequestWithContext(ctx, "POST", requestURL, bytes.NewReader(jsonData))
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Azure uses api-key header instead of Authorization: Bearer
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
if p.apiKey != "" {
|
||||||
|
req.Header.Set("Api-Key", p.apiKey)
|
||||||
|
}
|
||||||
|
|
||||||
|
resp, err := p.httpClient.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to send request: %w", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
|
||||||
|
if resp.StatusCode != http.StatusOK {
|
||||||
|
return nil, common.HandleErrorResponse(resp, p.apiBase)
|
||||||
|
}
|
||||||
|
|
||||||
|
return common.ReadAndParseResponse(resp, p.apiBase)
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetDefaultModel returns an empty string as Azure deployments are user-configured.
|
||||||
|
func (p *Provider) GetDefaultModel() string {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
232
pkg/providers/azure/provider_test.go
Normal file
232
pkg/providers/azure/provider_test.go
Normal file
|
|
@ -0,0 +1,232 @@
|
||||||
|
package azure
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// writeValidResponse writes a minimal valid Azure OpenAI chat completion response.
|
||||||
|
func writeValidResponse(w http.ResponseWriter) {
|
||||||
|
resp := map[string]any{
|
||||||
|
"choices": []map[string]any{
|
||||||
|
{
|
||||||
|
"message": map[string]any{"content": "ok"},
|
||||||
|
"finish_reason": "stop",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
json.NewEncoder(w).Encode(resp)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProviderChat_AzureURLConstruction(t *testing.T) {
|
||||||
|
var capturedPath string
|
||||||
|
var capturedAPIVersion string
|
||||||
|
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
capturedPath = r.URL.Path
|
||||||
|
capturedAPIVersion = r.URL.Query().Get("api-version")
|
||||||
|
writeValidResponse(w)
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
p := NewProvider("test-key", server.URL, "")
|
||||||
|
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "my-gpt5-deployment", nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Chat() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
wantPath := "/openai/deployments/my-gpt5-deployment/chat/completions"
|
||||||
|
if capturedPath != wantPath {
|
||||||
|
t.Errorf("URL path = %q, want %q", capturedPath, wantPath)
|
||||||
|
}
|
||||||
|
if capturedAPIVersion != azureAPIVersion {
|
||||||
|
t.Errorf("api-version = %q, want %q", capturedAPIVersion, azureAPIVersion)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProviderChat_AzureAuthHeader(t *testing.T) {
|
||||||
|
var capturedAPIKey string
|
||||||
|
var capturedAuth string
|
||||||
|
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
capturedAPIKey = r.Header.Get("Api-Key")
|
||||||
|
capturedAuth = r.Header.Get("Authorization")
|
||||||
|
writeValidResponse(w)
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
p := NewProvider("test-azure-key", server.URL, "")
|
||||||
|
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Chat() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if capturedAPIKey != "test-azure-key" {
|
||||||
|
t.Errorf("api-key header = %q, want %q", capturedAPIKey, "test-azure-key")
|
||||||
|
}
|
||||||
|
if capturedAuth != "" {
|
||||||
|
t.Errorf("Authorization header should be empty, got %q", capturedAuth)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProviderChat_AzureOmitsModelFromBody(t *testing.T) {
|
||||||
|
var requestBody map[string]any
|
||||||
|
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
json.NewDecoder(r.Body).Decode(&requestBody)
|
||||||
|
writeValidResponse(w)
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
p := NewProvider("test-key", server.URL, "")
|
||||||
|
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Chat() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if _, exists := requestBody["model"]; exists {
|
||||||
|
t.Error("request body should not contain 'model' field for Azure OpenAI")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProviderChat_AzureUsesMaxCompletionTokens(t *testing.T) {
|
||||||
|
var requestBody map[string]any
|
||||||
|
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
json.NewDecoder(r.Body).Decode(&requestBody)
|
||||||
|
writeValidResponse(w)
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
p := NewProvider("test-key", server.URL, "")
|
||||||
|
_, err := p.Chat(
|
||||||
|
t.Context(),
|
||||||
|
[]Message{{Role: "user", Content: "hi"}},
|
||||||
|
nil,
|
||||||
|
"deployment",
|
||||||
|
map[string]any{"max_tokens": 2048},
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Chat() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if _, exists := requestBody["max_completion_tokens"]; !exists {
|
||||||
|
t.Error("request body should contain 'max_completion_tokens'")
|
||||||
|
}
|
||||||
|
if _, exists := requestBody["max_tokens"]; exists {
|
||||||
|
t.Error("request body should not contain 'max_tokens'")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProviderChat_AzureHTTPError(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
http.Error(w, `{"error":"unauthorized"}`, http.StatusUnauthorized)
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
p := NewProvider("bad-key", server.URL, "")
|
||||||
|
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil)
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error, got nil")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProviderChat_AzureParseToolCalls(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
resp := map[string]any{
|
||||||
|
"choices": []map[string]any{
|
||||||
|
{
|
||||||
|
"message": map[string]any{
|
||||||
|
"content": "",
|
||||||
|
"tool_calls": []map[string]any{
|
||||||
|
{
|
||||||
|
"id": "call_1",
|
||||||
|
"type": "function",
|
||||||
|
"function": map[string]any{
|
||||||
|
"name": "get_weather",
|
||||||
|
"arguments": `{"city":"Seattle"}`,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"finish_reason": "tool_calls",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
json.NewEncoder(w).Encode(resp)
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
p := NewProvider("test-key", server.URL, "")
|
||||||
|
out, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "weather?"}}, nil, "deployment", nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Chat() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(out.ToolCalls) != 1 {
|
||||||
|
t.Fatalf("len(ToolCalls) = %d, want 1", len(out.ToolCalls))
|
||||||
|
}
|
||||||
|
if out.ToolCalls[0].Name != "get_weather" {
|
||||||
|
t.Errorf("ToolCalls[0].Name = %q, want %q", out.ToolCalls[0].Name, "get_weather")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProvider_AzureEmptyAPIBase(t *testing.T) {
|
||||||
|
p := NewProvider("test-key", "", "")
|
||||||
|
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil)
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for empty API base")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProvider_AzureRequestTimeoutDefault(t *testing.T) {
|
||||||
|
p := NewProvider("test-key", "https://example.com", "")
|
||||||
|
if p.httpClient.Timeout != defaultRequestTimeout {
|
||||||
|
t.Errorf("timeout = %v, want %v", p.httpClient.Timeout, defaultRequestTimeout)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProvider_AzureRequestTimeoutOverride(t *testing.T) {
|
||||||
|
p := NewProvider("test-key", "https://example.com", "", WithRequestTimeout(300*time.Second))
|
||||||
|
if p.httpClient.Timeout != 300*time.Second {
|
||||||
|
t.Errorf("timeout = %v, want %v", p.httpClient.Timeout, 300*time.Second)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProvider_AzureNewProviderWithTimeout(t *testing.T) {
|
||||||
|
p := NewProviderWithTimeout("test-key", "https://example.com", "", 180)
|
||||||
|
if p.httpClient.Timeout != 180*time.Second {
|
||||||
|
t.Errorf("timeout = %v, want %v", p.httpClient.Timeout, 180*time.Second)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProviderChat_AzureDeploymentNameEscaped(t *testing.T) {
|
||||||
|
var capturedPath string
|
||||||
|
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
capturedPath = r.URL.RawPath // use RawPath to see percent-encoding
|
||||||
|
if capturedPath == "" {
|
||||||
|
capturedPath = r.URL.Path
|
||||||
|
}
|
||||||
|
writeValidResponse(w)
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
p := NewProvider("test-key", server.URL, "")
|
||||||
|
|
||||||
|
// Deployment name with characters that could cause path injection
|
||||||
|
_, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "my deploy/../../admin", nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Chat() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// The slash and special chars in the deployment name must be escaped, not treated as path separators
|
||||||
|
if capturedPath == "/openai/deployments/my deploy/../../admin/chat/completions" {
|
||||||
|
t.Fatal("deployment name was interpolated without escaping — path injection possible")
|
||||||
|
}
|
||||||
|
}
|
||||||
380
pkg/providers/common/common.go
Normal file
380
pkg/providers/common/common.go
Normal file
|
|
@ -0,0 +1,380 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
// Package common provides shared utilities used by multiple LLM provider
|
||||||
|
// implementations (openai_compat, azure, etc.).
|
||||||
|
package common
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bufio"
|
||||||
|
"bytes"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"log"
|
||||||
|
"net/http"
|
||||||
|
"net/url"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/providers/protocoltypes"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Re-export protocol types used across providers.
|
||||||
|
type (
|
||||||
|
ToolCall = protocoltypes.ToolCall
|
||||||
|
FunctionCall = protocoltypes.FunctionCall
|
||||||
|
LLMResponse = protocoltypes.LLMResponse
|
||||||
|
UsageInfo = protocoltypes.UsageInfo
|
||||||
|
Message = protocoltypes.Message
|
||||||
|
ToolDefinition = protocoltypes.ToolDefinition
|
||||||
|
ToolFunctionDefinition = protocoltypes.ToolFunctionDefinition
|
||||||
|
ExtraContent = protocoltypes.ExtraContent
|
||||||
|
GoogleExtra = protocoltypes.GoogleExtra
|
||||||
|
ReasoningDetail = protocoltypes.ReasoningDetail
|
||||||
|
)
|
||||||
|
|
||||||
|
const DefaultRequestTimeout = 120 * time.Second
|
||||||
|
|
||||||
|
// NewHTTPClient creates an *http.Client with an optional proxy and the default timeout.
|
||||||
|
func NewHTTPClient(proxy string) *http.Client {
|
||||||
|
client := &http.Client{
|
||||||
|
Timeout: DefaultRequestTimeout,
|
||||||
|
}
|
||||||
|
if proxy != "" {
|
||||||
|
parsed, err := url.Parse(proxy)
|
||||||
|
if err == nil {
|
||||||
|
// Preserve http.DefaultTransport settings (TLS, HTTP/2, timeouts, etc.)
|
||||||
|
if base, ok := http.DefaultTransport.(*http.Transport); ok {
|
||||||
|
tr := base.Clone()
|
||||||
|
tr.Proxy = http.ProxyURL(parsed)
|
||||||
|
client.Transport = tr
|
||||||
|
} else {
|
||||||
|
// Fallback: minimal transport if DefaultTransport is not *http.Transport.
|
||||||
|
client.Transport = &http.Transport{
|
||||||
|
Proxy: http.ProxyURL(parsed),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
log.Printf("common: invalid proxy URL %q: %v", proxy, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return client
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Message serialization ---
|
||||||
|
|
||||||
|
// openaiMessage is the wire-format message for OpenAI-compatible APIs.
|
||||||
|
// It mirrors protocoltypes.Message but omits SystemParts, which is an
|
||||||
|
// internal field that would be unknown to third-party endpoints.
|
||||||
|
type openaiMessage struct {
|
||||||
|
Role string `json:"role"`
|
||||||
|
Content string `json:"content"`
|
||||||
|
ReasoningContent string `json:"reasoning_content,omitempty"`
|
||||||
|
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
|
||||||
|
ToolCallID string `json:"tool_call_id,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// SerializeMessages converts internal Message structs to the OpenAI wire format.
|
||||||
|
// - Strips SystemParts (unknown to third-party endpoints)
|
||||||
|
// - Converts messages with Media to multipart content format (text + image_url parts)
|
||||||
|
// - Preserves ToolCallID, ToolCalls, and ReasoningContent for all messages
|
||||||
|
func SerializeMessages(messages []Message) []any {
|
||||||
|
out := make([]any, 0, len(messages))
|
||||||
|
for _, m := range messages {
|
||||||
|
if len(m.Media) == 0 {
|
||||||
|
out = append(out, openaiMessage{
|
||||||
|
Role: m.Role,
|
||||||
|
Content: m.Content,
|
||||||
|
ReasoningContent: m.ReasoningContent,
|
||||||
|
ToolCalls: m.ToolCalls,
|
||||||
|
ToolCallID: m.ToolCallID,
|
||||||
|
})
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
// Multipart content format for messages with media
|
||||||
|
parts := make([]map[string]any, 0, 1+len(m.Media))
|
||||||
|
if m.Content != "" {
|
||||||
|
parts = append(parts, map[string]any{
|
||||||
|
"type": "text",
|
||||||
|
"text": m.Content,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
for _, mediaURL := range m.Media {
|
||||||
|
if strings.HasPrefix(mediaURL, "data:image/") {
|
||||||
|
parts = append(parts, map[string]any{
|
||||||
|
"type": "image_url",
|
||||||
|
"image_url": map[string]any{
|
||||||
|
"url": mediaURL,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
msg := map[string]any{
|
||||||
|
"role": m.Role,
|
||||||
|
"content": parts,
|
||||||
|
}
|
||||||
|
if m.ToolCallID != "" {
|
||||||
|
msg["tool_call_id"] = m.ToolCallID
|
||||||
|
}
|
||||||
|
if len(m.ToolCalls) > 0 {
|
||||||
|
msg["tool_calls"] = m.ToolCalls
|
||||||
|
}
|
||||||
|
if m.ReasoningContent != "" {
|
||||||
|
msg["reasoning_content"] = m.ReasoningContent
|
||||||
|
}
|
||||||
|
out = append(out, msg)
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Response parsing ---
|
||||||
|
|
||||||
|
// ParseResponse parses a JSON chat completion response body into an LLMResponse.
|
||||||
|
func ParseResponse(body io.Reader) (*LLMResponse, error) {
|
||||||
|
var apiResponse struct {
|
||||||
|
Choices []struct {
|
||||||
|
Message struct {
|
||||||
|
Content string `json:"content"`
|
||||||
|
ReasoningContent string `json:"reasoning_content"`
|
||||||
|
Reasoning string `json:"reasoning"`
|
||||||
|
ReasoningDetails []ReasoningDetail `json:"reasoning_details"`
|
||||||
|
ToolCalls []struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Type string `json:"type"`
|
||||||
|
Function *struct {
|
||||||
|
Name string `json:"name"`
|
||||||
|
Arguments json.RawMessage `json:"arguments"`
|
||||||
|
} `json:"function"`
|
||||||
|
ExtraContent *struct {
|
||||||
|
Google *struct {
|
||||||
|
ThoughtSignature string `json:"thought_signature"`
|
||||||
|
} `json:"google"`
|
||||||
|
} `json:"extra_content"`
|
||||||
|
} `json:"tool_calls"`
|
||||||
|
} `json:"message"`
|
||||||
|
FinishReason string `json:"finish_reason"`
|
||||||
|
} `json:"choices"`
|
||||||
|
Usage *UsageInfo `json:"usage"`
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := json.NewDecoder(body).Decode(&apiResponse); err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to decode response: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(apiResponse.Choices) == 0 {
|
||||||
|
return &LLMResponse{
|
||||||
|
Content: "",
|
||||||
|
FinishReason: "stop",
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
choice := apiResponse.Choices[0]
|
||||||
|
toolCalls := make([]ToolCall, 0, len(choice.Message.ToolCalls))
|
||||||
|
for _, tc := range choice.Message.ToolCalls {
|
||||||
|
arguments := make(map[string]any)
|
||||||
|
name := ""
|
||||||
|
|
||||||
|
// Extract thought_signature from Gemini/Google-specific extra content
|
||||||
|
thoughtSignature := ""
|
||||||
|
if tc.ExtraContent != nil && tc.ExtraContent.Google != nil {
|
||||||
|
thoughtSignature = tc.ExtraContent.Google.ThoughtSignature
|
||||||
|
}
|
||||||
|
|
||||||
|
if tc.Function != nil {
|
||||||
|
name = tc.Function.Name
|
||||||
|
arguments = DecodeToolCallArguments(tc.Function.Arguments, name)
|
||||||
|
}
|
||||||
|
|
||||||
|
toolCall := ToolCall{
|
||||||
|
ID: tc.ID,
|
||||||
|
Name: name,
|
||||||
|
Arguments: arguments,
|
||||||
|
ThoughtSignature: thoughtSignature,
|
||||||
|
}
|
||||||
|
|
||||||
|
if thoughtSignature != "" {
|
||||||
|
toolCall.ExtraContent = &ExtraContent{
|
||||||
|
Google: &GoogleExtra{
|
||||||
|
ThoughtSignature: thoughtSignature,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
toolCalls = append(toolCalls, toolCall)
|
||||||
|
}
|
||||||
|
|
||||||
|
return &LLMResponse{
|
||||||
|
Content: choice.Message.Content,
|
||||||
|
ReasoningContent: choice.Message.ReasoningContent,
|
||||||
|
Reasoning: choice.Message.Reasoning,
|
||||||
|
ReasoningDetails: choice.Message.ReasoningDetails,
|
||||||
|
ToolCalls: toolCalls,
|
||||||
|
FinishReason: choice.FinishReason,
|
||||||
|
Usage: apiResponse.Usage,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// DecodeToolCallArguments decodes a tool call's arguments from raw JSON.
|
||||||
|
func DecodeToolCallArguments(raw json.RawMessage, name string) map[string]any {
|
||||||
|
arguments := make(map[string]any)
|
||||||
|
raw = bytes.TrimSpace(raw)
|
||||||
|
if len(raw) == 0 || bytes.Equal(raw, []byte("null")) {
|
||||||
|
return arguments
|
||||||
|
}
|
||||||
|
|
||||||
|
var decoded any
|
||||||
|
if err := json.Unmarshal(raw, &decoded); err != nil {
|
||||||
|
log.Printf("common: failed to decode tool call arguments payload for %q: %v", name, err)
|
||||||
|
arguments["raw"] = string(raw)
|
||||||
|
return arguments
|
||||||
|
}
|
||||||
|
|
||||||
|
switch v := decoded.(type) {
|
||||||
|
case string:
|
||||||
|
if strings.TrimSpace(v) == "" {
|
||||||
|
return arguments
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal([]byte(v), &arguments); err != nil {
|
||||||
|
log.Printf("common: failed to decode tool call arguments for %q: %v", name, err)
|
||||||
|
arguments["raw"] = v
|
||||||
|
}
|
||||||
|
return arguments
|
||||||
|
case map[string]any:
|
||||||
|
return v
|
||||||
|
default:
|
||||||
|
log.Printf("common: unsupported tool call arguments type for %q: %T", name, decoded)
|
||||||
|
arguments["raw"] = string(raw)
|
||||||
|
return arguments
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- HTTP response helpers ---
|
||||||
|
|
||||||
|
// HandleErrorResponse reads a non-200 response body and returns an appropriate error.
|
||||||
|
func HandleErrorResponse(resp *http.Response, apiBase string) error {
|
||||||
|
contentType := resp.Header.Get("Content-Type")
|
||||||
|
body, readErr := io.ReadAll(io.LimitReader(resp.Body, 256))
|
||||||
|
if readErr != nil {
|
||||||
|
return fmt.Errorf("failed to read response: %w", readErr)
|
||||||
|
}
|
||||||
|
if LooksLikeHTML(body, contentType) {
|
||||||
|
return WrapHTMLResponseError(resp.StatusCode, body, contentType, apiBase)
|
||||||
|
}
|
||||||
|
return fmt.Errorf(
|
||||||
|
"API request failed:\n Status: %d\n Body: %s",
|
||||||
|
resp.StatusCode,
|
||||||
|
ResponsePreview(body, 128),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ReadAndParseResponse peeks at the response body to detect HTML errors,
|
||||||
|
// then parses the JSON response into an LLMResponse.
|
||||||
|
func ReadAndParseResponse(resp *http.Response, apiBase string) (*LLMResponse, error) {
|
||||||
|
contentType := resp.Header.Get("Content-Type")
|
||||||
|
reader := bufio.NewReader(resp.Body)
|
||||||
|
prefix, err := reader.Peek(256)
|
||||||
|
if err != nil && err != io.EOF && err != bufio.ErrBufferFull {
|
||||||
|
return nil, fmt.Errorf("failed to inspect response: %w", err)
|
||||||
|
}
|
||||||
|
if LooksLikeHTML(prefix, contentType) {
|
||||||
|
return nil, WrapHTMLResponseError(resp.StatusCode, prefix, contentType, apiBase)
|
||||||
|
}
|
||||||
|
out, err := ParseResponse(reader)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to parse JSON response: %w", err)
|
||||||
|
}
|
||||||
|
return out, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// LooksLikeHTML checks if the response body appears to be HTML.
|
||||||
|
func LooksLikeHTML(body []byte, contentType string) bool {
|
||||||
|
contentType = strings.ToLower(strings.TrimSpace(contentType))
|
||||||
|
if strings.Contains(contentType, "text/html") || strings.Contains(contentType, "application/xhtml+xml") {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
prefix := bytes.ToLower(leadingTrimmedPrefix(body, 128))
|
||||||
|
return bytes.HasPrefix(prefix, []byte("<!doctype html")) ||
|
||||||
|
bytes.HasPrefix(prefix, []byte("<html")) ||
|
||||||
|
bytes.HasPrefix(prefix, []byte("<head")) ||
|
||||||
|
bytes.HasPrefix(prefix, []byte("<body"))
|
||||||
|
}
|
||||||
|
|
||||||
|
// WrapHTMLResponseError creates a descriptive error for HTML responses.
|
||||||
|
func WrapHTMLResponseError(statusCode int, body []byte, contentType, apiBase string) error {
|
||||||
|
respPreview := ResponsePreview(body, 128)
|
||||||
|
return fmt.Errorf(
|
||||||
|
"API request failed: %s returned HTML instead of JSON (content-type: %s); check api_base or proxy configuration.\n Status: %d\n Body: %s",
|
||||||
|
apiBase,
|
||||||
|
contentType,
|
||||||
|
statusCode,
|
||||||
|
respPreview,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ResponsePreview returns a truncated preview of response body for error messages.
|
||||||
|
func ResponsePreview(body []byte, maxLen int) string {
|
||||||
|
trimmed := bytes.TrimSpace(body)
|
||||||
|
if len(trimmed) == 0 {
|
||||||
|
return "<empty>"
|
||||||
|
}
|
||||||
|
if len(trimmed) <= maxLen {
|
||||||
|
return string(trimmed)
|
||||||
|
}
|
||||||
|
return string(trimmed[:maxLen]) + "..."
|
||||||
|
}
|
||||||
|
|
||||||
|
func leadingTrimmedPrefix(body []byte, maxLen int) []byte {
|
||||||
|
i := 0
|
||||||
|
for i < len(body) {
|
||||||
|
switch body[i] {
|
||||||
|
case ' ', '\t', '\n', '\r', '\f', '\v':
|
||||||
|
i++
|
||||||
|
default:
|
||||||
|
end := i + maxLen
|
||||||
|
if end > len(body) {
|
||||||
|
end = len(body)
|
||||||
|
}
|
||||||
|
return body[i:end]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Numeric helpers ---
|
||||||
|
|
||||||
|
// AsInt converts various numeric types to int.
|
||||||
|
func AsInt(v any) (int, bool) {
|
||||||
|
switch val := v.(type) {
|
||||||
|
case int:
|
||||||
|
return val, true
|
||||||
|
case int64:
|
||||||
|
return int(val), true
|
||||||
|
case float64:
|
||||||
|
return int(val), true
|
||||||
|
case float32:
|
||||||
|
return int(val), true
|
||||||
|
default:
|
||||||
|
return 0, false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// AsFloat converts various numeric types to float64.
|
||||||
|
func AsFloat(v any) (float64, bool) {
|
||||||
|
switch val := v.(type) {
|
||||||
|
case float64:
|
||||||
|
return val, true
|
||||||
|
case float32:
|
||||||
|
return float64(val), true
|
||||||
|
case int:
|
||||||
|
return float64(val), true
|
||||||
|
case int64:
|
||||||
|
return float64(val), true
|
||||||
|
default:
|
||||||
|
return 0, false
|
||||||
|
}
|
||||||
|
}
|
||||||
558
pkg/providers/common/common_test.go
Normal file
558
pkg/providers/common/common_test.go
Normal file
|
|
@ -0,0 +1,558 @@
|
||||||
|
package common
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"net/url"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/providers/protocoltypes"
|
||||||
|
)
|
||||||
|
|
||||||
|
// --- NewHTTPClient tests ---
|
||||||
|
|
||||||
|
func TestNewHTTPClient_DefaultTimeout(t *testing.T) {
|
||||||
|
client := NewHTTPClient("")
|
||||||
|
if client.Timeout != DefaultRequestTimeout {
|
||||||
|
t.Errorf("timeout = %v, want %v", client.Timeout, DefaultRequestTimeout)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestNewHTTPClient_WithProxy(t *testing.T) {
|
||||||
|
client := NewHTTPClient("http://127.0.0.1:8080")
|
||||||
|
transport, ok := client.Transport.(*http.Transport)
|
||||||
|
if !ok || transport == nil {
|
||||||
|
t.Fatalf("expected http.Transport with proxy, got %T", client.Transport)
|
||||||
|
}
|
||||||
|
req := &http.Request{URL: &url.URL{Scheme: "https", Host: "api.example.com"}}
|
||||||
|
gotProxy, err := transport.Proxy(req)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("proxy function error: %v", err)
|
||||||
|
}
|
||||||
|
if gotProxy == nil || gotProxy.String() != "http://127.0.0.1:8080" {
|
||||||
|
t.Errorf("proxy = %v, want http://127.0.0.1:8080", gotProxy)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestNewHTTPClient_NoProxy(t *testing.T) {
|
||||||
|
client := NewHTTPClient("")
|
||||||
|
if client.Transport != nil {
|
||||||
|
t.Errorf("expected nil transport without proxy, got %T", client.Transport)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestNewHTTPClient_InvalidProxy(t *testing.T) {
|
||||||
|
// Should not panic, just log and return client without proxy
|
||||||
|
client := NewHTTPClient("://bad-url")
|
||||||
|
if client == nil {
|
||||||
|
t.Fatal("expected non-nil client even with invalid proxy")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- SerializeMessages tests ---
|
||||||
|
|
||||||
|
func TestSerializeMessages_PlainText(t *testing.T) {
|
||||||
|
messages := []Message{
|
||||||
|
{Role: "user", Content: "hello"},
|
||||||
|
{Role: "assistant", Content: "hi", ReasoningContent: "thinking..."},
|
||||||
|
}
|
||||||
|
result := SerializeMessages(messages)
|
||||||
|
|
||||||
|
data, _ := json.Marshal(result)
|
||||||
|
var msgs []map[string]any
|
||||||
|
json.Unmarshal(data, &msgs)
|
||||||
|
|
||||||
|
if msgs[0]["content"] != "hello" {
|
||||||
|
t.Errorf("expected plain string content, got %v", msgs[0]["content"])
|
||||||
|
}
|
||||||
|
if msgs[1]["reasoning_content"] != "thinking..." {
|
||||||
|
t.Errorf("reasoning_content not preserved, got %v", msgs[1]["reasoning_content"])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSerializeMessages_WithMedia(t *testing.T) {
|
||||||
|
messages := []Message{
|
||||||
|
{Role: "user", Content: "describe this", Media: []string{"data:image/png;base64,abc123"}},
|
||||||
|
}
|
||||||
|
result := SerializeMessages(messages)
|
||||||
|
|
||||||
|
data, _ := json.Marshal(result)
|
||||||
|
var msgs []map[string]any
|
||||||
|
json.Unmarshal(data, &msgs)
|
||||||
|
|
||||||
|
content, ok := msgs[0]["content"].([]any)
|
||||||
|
if !ok {
|
||||||
|
t.Fatalf("expected array content for media message, got %T", msgs[0]["content"])
|
||||||
|
}
|
||||||
|
if len(content) != 2 {
|
||||||
|
t.Fatalf("expected 2 content parts, got %d", len(content))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSerializeMessages_MediaWithToolCallID(t *testing.T) {
|
||||||
|
messages := []Message{
|
||||||
|
{Role: "tool", Content: "result", Media: []string{"data:image/png;base64,xyz"}, ToolCallID: "call_1"},
|
||||||
|
}
|
||||||
|
result := SerializeMessages(messages)
|
||||||
|
|
||||||
|
data, _ := json.Marshal(result)
|
||||||
|
var msgs []map[string]any
|
||||||
|
json.Unmarshal(data, &msgs)
|
||||||
|
|
||||||
|
if msgs[0]["tool_call_id"] != "call_1" {
|
||||||
|
t.Errorf("tool_call_id not preserved, got %v", msgs[0]["tool_call_id"])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSerializeMessages_StripsSystemParts(t *testing.T) {
|
||||||
|
messages := []Message{
|
||||||
|
{
|
||||||
|
Role: "system",
|
||||||
|
Content: "you are helpful",
|
||||||
|
SystemParts: []protocoltypes.ContentBlock{
|
||||||
|
{Type: "text", Text: "you are helpful"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
result := SerializeMessages(messages)
|
||||||
|
|
||||||
|
data, _ := json.Marshal(result)
|
||||||
|
if strings.Contains(string(data), "system_parts") {
|
||||||
|
t.Error("system_parts should not appear in serialized output")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- ParseResponse tests ---
|
||||||
|
|
||||||
|
func TestParseResponse_BasicContent(t *testing.T) {
|
||||||
|
body := `{"choices":[{"message":{"content":"hello world"},"finish_reason":"stop"}]}`
|
||||||
|
out, err := ParseResponse(strings.NewReader(body))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ParseResponse() error = %v", err)
|
||||||
|
}
|
||||||
|
if out.Content != "hello world" {
|
||||||
|
t.Errorf("Content = %q, want %q", out.Content, "hello world")
|
||||||
|
}
|
||||||
|
if out.FinishReason != "stop" {
|
||||||
|
t.Errorf("FinishReason = %q, want %q", out.FinishReason, "stop")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParseResponse_EmptyChoices(t *testing.T) {
|
||||||
|
body := `{"choices":[]}`
|
||||||
|
out, err := ParseResponse(strings.NewReader(body))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ParseResponse() error = %v", err)
|
||||||
|
}
|
||||||
|
if out.Content != "" {
|
||||||
|
t.Errorf("Content = %q, want empty", out.Content)
|
||||||
|
}
|
||||||
|
if out.FinishReason != "stop" {
|
||||||
|
t.Errorf("FinishReason = %q, want %q", out.FinishReason, "stop")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParseResponse_WithToolCalls(t *testing.T) {
|
||||||
|
body := `{"choices":[{"message":{"content":"","tool_calls":[{"id":"call_1","type":"function","function":{"name":"get_weather","arguments":"{\"city\":\"SF\"}"}}]},"finish_reason":"tool_calls"}]}`
|
||||||
|
out, err := ParseResponse(strings.NewReader(body))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ParseResponse() error = %v", err)
|
||||||
|
}
|
||||||
|
if len(out.ToolCalls) != 1 {
|
||||||
|
t.Fatalf("len(ToolCalls) = %d, want 1", len(out.ToolCalls))
|
||||||
|
}
|
||||||
|
if out.ToolCalls[0].Name != "get_weather" {
|
||||||
|
t.Errorf("ToolCalls[0].Name = %q, want %q", out.ToolCalls[0].Name, "get_weather")
|
||||||
|
}
|
||||||
|
if out.ToolCalls[0].Arguments["city"] != "SF" {
|
||||||
|
t.Errorf("ToolCalls[0].Arguments[city] = %v, want SF", out.ToolCalls[0].Arguments["city"])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParseResponse_WithUsage(t *testing.T) {
|
||||||
|
body := `{"choices":[{"message":{"content":"ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}`
|
||||||
|
out, err := ParseResponse(strings.NewReader(body))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ParseResponse() error = %v", err)
|
||||||
|
}
|
||||||
|
if out.Usage == nil {
|
||||||
|
t.Fatal("Usage is nil")
|
||||||
|
}
|
||||||
|
if out.Usage.PromptTokens != 10 {
|
||||||
|
t.Errorf("PromptTokens = %d, want 10", out.Usage.PromptTokens)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParseResponse_WithReasoningContent(t *testing.T) {
|
||||||
|
body := `{"choices":[{"message":{"content":"2","reasoning_content":"Let me think... 1+1=2"},"finish_reason":"stop"}]}`
|
||||||
|
out, err := ParseResponse(strings.NewReader(body))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ParseResponse() error = %v", err)
|
||||||
|
}
|
||||||
|
if out.ReasoningContent != "Let me think... 1+1=2" {
|
||||||
|
t.Errorf("ReasoningContent = %q, want %q", out.ReasoningContent, "Let me think... 1+1=2")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParseResponse_InvalidJSON(t *testing.T) {
|
||||||
|
_, err := ParseResponse(strings.NewReader("not json"))
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for invalid JSON")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- DecodeToolCallArguments tests ---
|
||||||
|
|
||||||
|
func TestDecodeToolCallArguments_ObjectJSON(t *testing.T) {
|
||||||
|
raw := json.RawMessage(`{"city":"Seattle","units":"metric"}`)
|
||||||
|
args := DecodeToolCallArguments(raw, "test")
|
||||||
|
if args["city"] != "Seattle" {
|
||||||
|
t.Errorf("city = %v, want Seattle", args["city"])
|
||||||
|
}
|
||||||
|
if args["units"] != "metric" {
|
||||||
|
t.Errorf("units = %v, want metric", args["units"])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDecodeToolCallArguments_StringJSON(t *testing.T) {
|
||||||
|
raw := json.RawMessage(`"{\"city\":\"SF\"}"`)
|
||||||
|
args := DecodeToolCallArguments(raw, "test")
|
||||||
|
if args["city"] != "SF" {
|
||||||
|
t.Errorf("city = %v, want SF", args["city"])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDecodeToolCallArguments_EmptyInput(t *testing.T) {
|
||||||
|
args := DecodeToolCallArguments(nil, "test")
|
||||||
|
if len(args) != 0 {
|
||||||
|
t.Errorf("expected empty map, got %v", args)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDecodeToolCallArguments_NullInput(t *testing.T) {
|
||||||
|
args := DecodeToolCallArguments(json.RawMessage(`null`), "test")
|
||||||
|
if len(args) != 0 {
|
||||||
|
t.Errorf("expected empty map, got %v", args)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDecodeToolCallArguments_InvalidJSON(t *testing.T) {
|
||||||
|
args := DecodeToolCallArguments(json.RawMessage(`not-json`), "test")
|
||||||
|
if _, ok := args["raw"]; !ok {
|
||||||
|
t.Error("expected 'raw' fallback key for invalid JSON")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDecodeToolCallArguments_EmptyStringJSON(t *testing.T) {
|
||||||
|
args := DecodeToolCallArguments(json.RawMessage(`" "`), "test")
|
||||||
|
if len(args) != 0 {
|
||||||
|
t.Errorf("expected empty map for whitespace string, got %v", args)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- HandleErrorResponse tests ---
|
||||||
|
|
||||||
|
func TestHandleErrorResponse_JSONError(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(http.StatusBadRequest)
|
||||||
|
w.Write([]byte(`{"error":"bad request"}`))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
resp, err := http.Get(server.URL)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("http.Get() error = %v", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
err = HandleErrorResponse(resp, server.URL)
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "400") {
|
||||||
|
t.Errorf("error should contain status code, got %v", err)
|
||||||
|
}
|
||||||
|
if strings.Contains(err.Error(), "HTML") {
|
||||||
|
t.Errorf("should not mention HTML for JSON error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandleErrorResponse_HTMLError(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/html")
|
||||||
|
w.WriteHeader(http.StatusBadGateway)
|
||||||
|
w.Write([]byte("<!DOCTYPE html><html><body>bad gateway</body></html>"))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
resp, err := http.Get(server.URL)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("http.Get() error = %v", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
err = HandleErrorResponse(resp, server.URL)
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "HTML instead of JSON") {
|
||||||
|
t.Errorf("expected HTML error message, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- ReadAndParseResponse tests ---
|
||||||
|
|
||||||
|
func TestReadAndParseResponse_ValidJSON(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.Write([]byte(`{"choices":[{"message":{"content":"ok"},"finish_reason":"stop"}]}`))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
resp, err := http.Get(server.URL)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("http.Get() error = %v", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
out, err := ReadAndParseResponse(resp, server.URL)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ReadAndParseResponse() error = %v", err)
|
||||||
|
}
|
||||||
|
if out.Content != "ok" {
|
||||||
|
t.Errorf("Content = %q, want %q", out.Content, "ok")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestReadAndParseResponse_HTMLResponse(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/html")
|
||||||
|
w.Write([]byte("<!DOCTYPE html><html><body>login page</body></html>"))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
resp, err := http.Get(server.URL)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("http.Get() error = %v", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
_, err = ReadAndParseResponse(resp, server.URL)
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for HTML response")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "HTML instead of JSON") {
|
||||||
|
t.Errorf("expected HTML error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- LooksLikeHTML tests ---
|
||||||
|
|
||||||
|
func TestLooksLikeHTML_ContentTypeHTML(t *testing.T) {
|
||||||
|
if !LooksLikeHTML(nil, "text/html; charset=utf-8") {
|
||||||
|
t.Error("expected true for text/html content type")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLooksLikeHTML_ContentTypeXHTML(t *testing.T) {
|
||||||
|
if !LooksLikeHTML(nil, "application/xhtml+xml") {
|
||||||
|
t.Error("expected true for xhtml content type")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLooksLikeHTML_BodyPrefix(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
body string
|
||||||
|
}{
|
||||||
|
{"doctype", "<!DOCTYPE html><html>"},
|
||||||
|
{"html tag", "<html><body>"},
|
||||||
|
{"head tag", "<head><title>"},
|
||||||
|
{"body tag", "<body>content"},
|
||||||
|
{"whitespace before", " \n\t<!DOCTYPE html>"},
|
||||||
|
}
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
if !LooksLikeHTML([]byte(tt.body), "application/json") {
|
||||||
|
t.Errorf("expected true for body %q", tt.body)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLooksLikeHTML_NotHTML(t *testing.T) {
|
||||||
|
if LooksLikeHTML([]byte(`{"error":"bad"}`), "application/json") {
|
||||||
|
t.Error("expected false for JSON body")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- ResponsePreview tests ---
|
||||||
|
|
||||||
|
func TestResponsePreview_Short(t *testing.T) {
|
||||||
|
got := ResponsePreview([]byte("hello"), 128)
|
||||||
|
if got != "hello" {
|
||||||
|
t.Errorf("got %q, want %q", got, "hello")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResponsePreview_Truncated(t *testing.T) {
|
||||||
|
body := strings.Repeat("a", 200)
|
||||||
|
got := ResponsePreview([]byte(body), 128)
|
||||||
|
if len(got) != 131 { // 128 + "..."
|
||||||
|
t.Errorf("len = %d, want 131", len(got))
|
||||||
|
}
|
||||||
|
if !strings.HasSuffix(got, "...") {
|
||||||
|
t.Error("expected ... suffix")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResponsePreview_Empty(t *testing.T) {
|
||||||
|
got := ResponsePreview([]byte(""), 128)
|
||||||
|
if got != "<empty>" {
|
||||||
|
t.Errorf("got %q, want %q", got, "<empty>")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResponsePreview_Whitespace(t *testing.T) {
|
||||||
|
got := ResponsePreview([]byte(" \n\t "), 128)
|
||||||
|
if got != "<empty>" {
|
||||||
|
t.Errorf("got %q, want %q for whitespace-only body", got, "<empty>")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- AsInt tests ---
|
||||||
|
|
||||||
|
func TestAsInt(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
val any
|
||||||
|
want int
|
||||||
|
ok bool
|
||||||
|
}{
|
||||||
|
{"int", 42, 42, true},
|
||||||
|
{"int64", int64(99), 99, true},
|
||||||
|
{"float64", float64(512), 512, true},
|
||||||
|
{"float32", float32(256), 256, true},
|
||||||
|
{"string", "nope", 0, false},
|
||||||
|
{"nil", nil, 0, false},
|
||||||
|
}
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
got, ok := AsInt(tt.val)
|
||||||
|
if ok != tt.ok || got != tt.want {
|
||||||
|
t.Errorf("AsInt(%v) = (%d, %v), want (%d, %v)", tt.val, got, ok, tt.want, tt.ok)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- AsFloat tests ---
|
||||||
|
|
||||||
|
func TestAsFloat(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
val any
|
||||||
|
want float64
|
||||||
|
ok bool
|
||||||
|
}{
|
||||||
|
{"float64", float64(0.7), 0.7, true},
|
||||||
|
{"float32", float32(0.5), float64(float32(0.5)), true},
|
||||||
|
{"int", 1, 1.0, true},
|
||||||
|
{"int64", int64(100), 100.0, true},
|
||||||
|
{"string", "nope", 0, false},
|
||||||
|
{"nil", nil, 0, false},
|
||||||
|
}
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
got, ok := AsFloat(tt.val)
|
||||||
|
if ok != tt.ok || got != tt.want {
|
||||||
|
t.Errorf("AsFloat(%v) = (%f, %v), want (%f, %v)", tt.val, got, ok, tt.want, tt.ok)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- WrapHTMLResponseError tests ---
|
||||||
|
|
||||||
|
func TestWrapHTMLResponseError(t *testing.T) {
|
||||||
|
err := WrapHTMLResponseError(502, []byte("<html>bad</html>"), "text/html", "https://api.example.com")
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error")
|
||||||
|
}
|
||||||
|
msg := err.Error()
|
||||||
|
if !strings.Contains(msg, "502") {
|
||||||
|
t.Errorf("expected status code in error, got %v", msg)
|
||||||
|
}
|
||||||
|
if !strings.Contains(msg, "https://api.example.com") {
|
||||||
|
t.Errorf("expected api base in error, got %v", msg)
|
||||||
|
}
|
||||||
|
if !strings.Contains(msg, "HTML instead of JSON") {
|
||||||
|
t.Errorf("expected HTML mention in error, got %v", msg)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- HandleErrorResponse with read failure ---
|
||||||
|
|
||||||
|
func TestHandleErrorResponse_EmptyBody(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(http.StatusInternalServerError)
|
||||||
|
// empty body
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
resp, err := http.Get(server.URL)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("http.Get() error = %v", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
err = HandleErrorResponse(resp, server.URL)
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "500") {
|
||||||
|
t.Errorf("expected status code, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- ReadAndParseResponse with invalid JSON ---
|
||||||
|
|
||||||
|
func TestReadAndParseResponse_InvalidJSON(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.Write([]byte("not valid json"))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
resp, err := http.Get(server.URL)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("http.Get() error = %v", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
_, err = ReadAndParseResponse(resp, server.URL)
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for invalid JSON")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- ParseResponse with thought_signature (Google/Gemini) ---
|
||||||
|
|
||||||
|
func TestParseResponse_WithThoughtSignature(t *testing.T) {
|
||||||
|
body := `{"choices":[{"message":{"content":"","tool_calls":[{"id":"call_1","type":"function","function":{"name":"test_tool","arguments":"{}"},"extra_content":{"google":{"thought_signature":"sig123"}}}]},"finish_reason":"tool_calls"}]}`
|
||||||
|
out, err := ParseResponse(strings.NewReader(body))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ParseResponse() error = %v", err)
|
||||||
|
}
|
||||||
|
if len(out.ToolCalls) != 1 {
|
||||||
|
t.Fatalf("len(ToolCalls) = %d, want 1", len(out.ToolCalls))
|
||||||
|
}
|
||||||
|
if out.ToolCalls[0].ThoughtSignature != "sig123" {
|
||||||
|
t.Errorf("ThoughtSignature = %q, want %q", out.ToolCalls[0].ThoughtSignature, "sig123")
|
||||||
|
}
|
||||||
|
if out.ToolCalls[0].ExtraContent == nil || out.ToolCalls[0].ExtraContent.Google == nil {
|
||||||
|
t.Fatal("ExtraContent.Google is nil")
|
||||||
|
}
|
||||||
|
if out.ToolCalls[0].ExtraContent.Google.ThoughtSignature != "sig123" {
|
||||||
|
t.Errorf("ExtraContent.Google.ThoughtSignature = %q, want %q",
|
||||||
|
out.ToolCalls[0].ExtraContent.Google.ThoughtSignature, "sig123")
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Add table
Reference in a new issue