fix: refactor shell streaming to use event bus pattern (#50)
This commit is contained in:
@@ -225,16 +225,6 @@ func (a *orchestratorAdapter) ResumeTaskSession(ctx context.Context, taskID, tas
|
||||
return err
|
||||
}
|
||||
|
||||
// shellSessionResumerAdapter adapts the orchestrator service to the shell handlers' SessionResumer interface
|
||||
type shellSessionResumerAdapter struct {
|
||||
svc *orchestrator.Service
|
||||
}
|
||||
|
||||
func (a *shellSessionResumerAdapter) ResumeTaskSession(ctx context.Context, taskID, taskSessionID string) error {
|
||||
_, err := a.svc.ResumeTaskSession(ctx, taskID, taskSessionID)
|
||||
return err
|
||||
}
|
||||
|
||||
// messageCreatorAdapter adapts the task service to the orchestrator.MessageCreator interface
|
||||
type messageCreatorAdapter struct {
|
||||
svc *taskservice.Service
|
||||
|
||||
@@ -55,10 +55,7 @@ func provideGateway(
|
||||
workspaceFileHandlers.RegisterHandlers(gateway.Dispatcher)
|
||||
|
||||
shellHandlers := agenthandlers.NewShellHandlers(lifecycleMgr, log)
|
||||
shellHandlers.SetHub(gateway.Hub)
|
||||
shellHandlers.SetSessionResumer(&shellSessionResumerAdapter{svc: orchestratorSvc})
|
||||
shellHandlers.RegisterHandlers(gateway.Dispatcher)
|
||||
lifecycleMgr.SetShellStreamStarter(shellHandlers)
|
||||
}
|
||||
|
||||
go gateway.Hub.Run(ctx)
|
||||
|
||||
@@ -561,6 +561,62 @@ func main() {
|
||||
log.Info("Subscribed to file change events for WebSocket broadcasting")
|
||||
}
|
||||
|
||||
// Subscribe to shell output events and broadcast to WebSocket subscribers
|
||||
_, err = eventBus.Subscribe(events.BuildShellOutputWildcardSubject(), func(ctx context.Context, event *bus.Event) error {
|
||||
// Extract session_id from the event data
|
||||
var taskSessionID string
|
||||
switch data := event.Data.(type) {
|
||||
case lifecycle.ShellOutputEventPayload:
|
||||
taskSessionID = data.SessionID
|
||||
case *lifecycle.ShellOutputEventPayload:
|
||||
taskSessionID = data.SessionID
|
||||
case map[string]interface{}:
|
||||
taskSessionID, _ = data["session_id"].(string)
|
||||
}
|
||||
if taskSessionID == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Broadcast shell output to session subscribers
|
||||
notification, _ := ws.NewNotification(ws.ActionSessionShellOutput, event.Data)
|
||||
gateway.Hub.BroadcastToSession(taskSessionID, notification)
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
log.Error("Failed to subscribe to shell output events", zap.Error(err))
|
||||
} else {
|
||||
log.Info("Subscribed to shell output events for WebSocket broadcasting")
|
||||
}
|
||||
|
||||
// Subscribe to shell exit events and broadcast to WebSocket subscribers
|
||||
// Note: Frontend expects shell exit via session.shell.output with type: "exit"
|
||||
_, err = eventBus.Subscribe(events.BuildShellExitWildcardSubject(), func(ctx context.Context, event *bus.Event) error {
|
||||
// Extract session_id from the event data
|
||||
var taskSessionID string
|
||||
switch data := event.Data.(type) {
|
||||
case lifecycle.ShellExitEventPayload:
|
||||
taskSessionID = data.SessionID
|
||||
case *lifecycle.ShellExitEventPayload:
|
||||
taskSessionID = data.SessionID
|
||||
case map[string]interface{}:
|
||||
taskSessionID, _ = data["session_id"].(string)
|
||||
}
|
||||
if taskSessionID == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Broadcast shell exit to session subscribers via session.shell.output action
|
||||
// Frontend handles exit events via type: "exit" in the same handler
|
||||
notification, _ := ws.NewNotification(ws.ActionSessionShellOutput, event.Data)
|
||||
gateway.Hub.BroadcastToSession(taskSessionID, notification)
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
log.Error("Failed to subscribe to shell exit events", zap.Error(err))
|
||||
} else {
|
||||
log.Info("Subscribed to shell exit events for WebSocket broadcasting")
|
||||
}
|
||||
|
||||
// ============================================
|
||||
// HTTP SERVER (WebSocket + HTTP endpoints)
|
||||
// ============================================
|
||||
|
||||
@@ -4,58 +4,29 @@ package handlers
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/kandev/kandev/internal/agent/lifecycle"
|
||||
agentctl "github.com/kandev/kandev/internal/agentctl/client"
|
||||
"github.com/kandev/kandev/internal/common/logger"
|
||||
ws "github.com/kandev/kandev/pkg/websocket"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// SessionResumer is an interface for resuming task sessions
|
||||
type SessionResumer interface {
|
||||
ResumeTaskSession(ctx context.Context, taskID, taskSessionID string) error
|
||||
}
|
||||
|
||||
// ShellHandlers provides WebSocket handlers for shell terminal operations
|
||||
// ShellHandlers provides WebSocket handlers for shell terminal operations.
|
||||
// Shell output is streamed via the lifecycle manager and event bus.
|
||||
// This handler provides shell.status, shell.subscribe (for buffer), and shell.input.
|
||||
type ShellHandlers struct {
|
||||
lifecycleMgr *lifecycle.Manager
|
||||
logger *logger.Logger
|
||||
hub ShellOutputBroadcaster
|
||||
sessionResumer SessionResumer
|
||||
|
||||
// Track active shell streams per session
|
||||
activeStreams map[string]context.CancelFunc
|
||||
mu sync.RWMutex
|
||||
|
||||
// Input channels for sending input to shell streams
|
||||
inputChannels map[string]chan<- agentctl.ShellMessage
|
||||
inputMu sync.RWMutex
|
||||
lifecycleMgr *lifecycle.Manager
|
||||
logger *logger.Logger
|
||||
}
|
||||
|
||||
// NewShellHandlers creates a new ShellHandlers instance
|
||||
func NewShellHandlers(lifecycleMgr *lifecycle.Manager, log *logger.Logger) *ShellHandlers {
|
||||
return &ShellHandlers{
|
||||
lifecycleMgr: lifecycleMgr,
|
||||
logger: log.WithFields(zap.String("component", "shell_handlers")),
|
||||
activeStreams: make(map[string]context.CancelFunc),
|
||||
inputChannels: make(map[string]chan<- agentctl.ShellMessage),
|
||||
lifecycleMgr: lifecycleMgr,
|
||||
logger: log.WithFields(zap.String("component", "shell_handlers")),
|
||||
}
|
||||
}
|
||||
|
||||
// SetHub sets the hub for broadcasting shell output
|
||||
func (h *ShellHandlers) SetHub(hub ShellOutputBroadcaster) {
|
||||
h.hub = hub
|
||||
}
|
||||
|
||||
// SetSessionResumer sets the session resumer for auto-resuming sessions on shell subscribe
|
||||
func (h *ShellHandlers) SetSessionResumer(resumer SessionResumer) {
|
||||
h.sessionResumer = resumer
|
||||
}
|
||||
|
||||
|
||||
// RegisterHandlers registers shell handlers with the WebSocket dispatcher
|
||||
func (h *ShellHandlers) RegisterHandlers(d *ws.Dispatcher) {
|
||||
d.RegisterFunc(ws.ActionShellStatus, h.wsShellStatus)
|
||||
@@ -126,7 +97,9 @@ type ShellSubscribeRequest struct {
|
||||
SessionID string `json:"session_id"`
|
||||
}
|
||||
|
||||
// wsShellSubscribe subscribes to shell output for a session and starts the shell stream if needed
|
||||
// wsShellSubscribe subscribes to shell output for a session.
|
||||
// Shell output is streamed via the event bus (lifecycle manager handles this).
|
||||
// This endpoint returns the buffered shell output for catchup.
|
||||
func (h *ShellHandlers) wsShellSubscribe(ctx context.Context, msg *ws.Message) (*ws.Message, error) {
|
||||
var req ShellSubscribeRequest
|
||||
if err := msg.ParsePayload(&req); err != nil {
|
||||
@@ -143,13 +116,9 @@ func (h *ShellHandlers) wsShellSubscribe(ctx context.Context, msg *ws.Message) (
|
||||
return nil, fmt.Errorf("no agent running for session %s", req.SessionID)
|
||||
}
|
||||
|
||||
// Start the shell stream (idempotent - does nothing if already running)
|
||||
if err := h.StartShellStream(ctx, req.SessionID); err != nil {
|
||||
return nil, fmt.Errorf("failed to start shell stream: %w", err)
|
||||
}
|
||||
|
||||
// Get buffered output to include in response
|
||||
// This ensures client gets current shell state without duplicate broadcasts
|
||||
// Shell output streaming is handled by the lifecycle manager via event bus
|
||||
buffer := ""
|
||||
if client := execution.GetAgentCtlClient(); client != nil {
|
||||
if b, err := client.ShellBuffer(ctx); err == nil {
|
||||
@@ -175,169 +144,23 @@ func (h *ShellHandlers) wsShellInput(ctx context.Context, msg *ws.Message) (*ws.
|
||||
return nil, fmt.Errorf("session_id is required")
|
||||
}
|
||||
|
||||
// Verify the agent execution exists for this session
|
||||
if _, ok := h.lifecycleMgr.GetExecutionBySessionID(req.SessionID); !ok {
|
||||
// Get the agent execution for this session
|
||||
execution, ok := h.lifecycleMgr.GetExecutionBySessionID(req.SessionID)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("no agent running for session %s", req.SessionID)
|
||||
}
|
||||
|
||||
// Send input via the shell stream
|
||||
if err := h.SendShellInput(req.SessionID, req.Data); err != nil {
|
||||
return nil, err
|
||||
// Use the workspace stream from AgentExecution (managed by lifecycle)
|
||||
workspaceStream := execution.GetWorkspaceStream()
|
||||
if workspaceStream == nil {
|
||||
return nil, fmt.Errorf("no workspace stream for session %s", req.SessionID)
|
||||
}
|
||||
|
||||
if err := workspaceStream.WriteShellInput(req.Data); err != nil {
|
||||
return nil, fmt.Errorf("failed to send shell input: %w", err)
|
||||
}
|
||||
|
||||
return ws.NewResponse(msg.ID, msg.Action, map[string]interface{}{
|
||||
"success": true,
|
||||
})
|
||||
}
|
||||
|
||||
// StartShellStream starts streaming shell output for a session using the configured hub.
|
||||
// This implements the lifecycle.ShellStreamStarter interface.
|
||||
func (h *ShellHandlers) StartShellStream(ctx context.Context, sessionID string) error {
|
||||
if h.hub == nil {
|
||||
return fmt.Errorf("hub not configured")
|
||||
}
|
||||
return h.StartShellStreamWithHub(ctx, sessionID, h.hub)
|
||||
}
|
||||
|
||||
// StartShellStreamWithHub starts streaming shell output for a session to a specific hub
|
||||
func (h *ShellHandlers) StartShellStreamWithHub(ctx context.Context, sessionID string, hub ShellOutputBroadcaster) error {
|
||||
h.mu.Lock()
|
||||
if _, exists := h.activeStreams[sessionID]; exists {
|
||||
h.mu.Unlock()
|
||||
return nil // Already streaming
|
||||
}
|
||||
|
||||
streamCtx, cancel := context.WithCancel(context.Background()) // Use background context so stream survives request
|
||||
h.activeStreams[sessionID] = cancel
|
||||
h.mu.Unlock()
|
||||
|
||||
// Use a channel to wait for stream setup to complete
|
||||
readyCh := make(chan error, 1)
|
||||
go h.runShellStreamWithReady(streamCtx, sessionID, hub, readyCh)
|
||||
|
||||
// Wait for stream to be ready (or fail)
|
||||
select {
|
||||
case err := <-readyCh:
|
||||
return err
|
||||
case <-ctx.Done():
|
||||
cancel()
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
// StopShellStream stops the shell stream for a session
|
||||
func (h *ShellHandlers) StopShellStream(sessionID string) {
|
||||
h.mu.Lock()
|
||||
if cancel, exists := h.activeStreams[sessionID]; exists {
|
||||
cancel()
|
||||
delete(h.activeStreams, sessionID)
|
||||
}
|
||||
h.mu.Unlock()
|
||||
}
|
||||
|
||||
// ShellOutputBroadcaster interface for broadcasting shell output
|
||||
type ShellOutputBroadcaster interface {
|
||||
BroadcastToSession(sessionID string, msg *ws.Message)
|
||||
}
|
||||
|
||||
// runShellStreamWithReady runs the shell output stream for a session and signals when ready
|
||||
func (h *ShellHandlers) runShellStreamWithReady(ctx context.Context, sessionID string, hub ShellOutputBroadcaster, readyCh chan<- error) {
|
||||
defer func() {
|
||||
h.mu.Lock()
|
||||
delete(h.activeStreams, sessionID)
|
||||
h.mu.Unlock()
|
||||
}()
|
||||
|
||||
execution, ok := h.lifecycleMgr.GetExecutionBySessionID(sessionID)
|
||||
if !ok {
|
||||
h.logger.Debug("no agent execution for shell stream", zap.String("session_id", sessionID))
|
||||
readyCh <- fmt.Errorf("no agent execution for session %s", sessionID)
|
||||
return
|
||||
}
|
||||
|
||||
client := execution.GetAgentCtlClient()
|
||||
if client == nil {
|
||||
h.logger.Debug("no client for shell stream", zap.String("session_id", sessionID))
|
||||
readyCh <- fmt.Errorf("no agent client for session %s", sessionID)
|
||||
return
|
||||
}
|
||||
|
||||
if err := client.WaitForReady(ctx, 10*time.Second); err != nil {
|
||||
h.logger.Error("agentctl not ready for shell stream", zap.String("session_id", sessionID), zap.Error(err))
|
||||
readyCh <- fmt.Errorf("agentctl not ready: %w", err)
|
||||
return
|
||||
}
|
||||
|
||||
outputCh, inputCh, err := client.StreamShell(ctx)
|
||||
if err != nil {
|
||||
h.logger.Error("failed to start shell stream", zap.String("session_id", sessionID), zap.Error(err))
|
||||
readyCh <- fmt.Errorf("failed to start shell stream: %w", err)
|
||||
return
|
||||
}
|
||||
|
||||
// Store input channel for sending input
|
||||
h.storeInputChannel(sessionID, inputCh)
|
||||
defer h.removeInputChannel(sessionID)
|
||||
|
||||
h.logger.Info("shell stream started", zap.String("session_id", sessionID))
|
||||
|
||||
// Signal that we're ready - input channel is now stored
|
||||
readyCh <- nil
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case msg, ok := <-outputCh:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
// Broadcast shell output to session subscribers
|
||||
notification, err := ws.NewNotification(ws.ActionShellOutput, map[string]interface{}{
|
||||
"session_id": sessionID,
|
||||
"type": msg.Type,
|
||||
"data": msg.Data,
|
||||
"code": msg.Code,
|
||||
})
|
||||
if err != nil {
|
||||
h.logger.Debug("failed to create shell output notification", zap.Error(err))
|
||||
continue
|
||||
}
|
||||
hub.BroadcastToSession(sessionID, notification)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
func (h *ShellHandlers) storeInputChannel(sessionID string, ch chan<- agentctl.ShellMessage) {
|
||||
h.inputMu.Lock()
|
||||
h.inputChannels[sessionID] = ch
|
||||
h.inputMu.Unlock()
|
||||
}
|
||||
|
||||
func (h *ShellHandlers) removeInputChannel(sessionID string) {
|
||||
h.inputMu.Lock()
|
||||
if ch, exists := h.inputChannels[sessionID]; exists {
|
||||
close(ch)
|
||||
delete(h.inputChannels, sessionID)
|
||||
}
|
||||
h.inputMu.Unlock()
|
||||
}
|
||||
|
||||
// SendShellInput sends input to an active shell stream
|
||||
func (h *ShellHandlers) SendShellInput(sessionID, data string) error {
|
||||
h.inputMu.RLock()
|
||||
ch, exists := h.inputChannels[sessionID]
|
||||
h.inputMu.RUnlock()
|
||||
|
||||
if !exists {
|
||||
return fmt.Errorf("no active shell stream for session %s", sessionID)
|
||||
}
|
||||
|
||||
select {
|
||||
case ch <- agentctl.ShellMessage{Type: "input", Data: data}:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("shell input channel full")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -91,12 +91,6 @@ func TestNewShellHandlers(t *testing.T) {
|
||||
if handlers.logger == nil {
|
||||
t.Error("expected non-nil logger")
|
||||
}
|
||||
if handlers.activeStreams == nil {
|
||||
t.Error("expected non-nil activeStreams map")
|
||||
}
|
||||
if handlers.inputChannels == nil {
|
||||
t.Error("expected non-nil inputChannels map")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisterHandlers(t *testing.T) {
|
||||
@@ -265,29 +259,6 @@ func TestWsShellInput_NoInstanceFound(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestSendShellInput_NoActiveStream(t *testing.T) {
|
||||
log := newTestLogger()
|
||||
handlers := NewShellHandlers(nil, log)
|
||||
|
||||
err := handlers.SendShellInput("non-existent-session", "test input")
|
||||
if err == nil {
|
||||
t.Error("expected error for non-existent stream")
|
||||
}
|
||||
expectedErr := "no active shell stream for session non-existent-session"
|
||||
if err.Error() != expectedErr {
|
||||
t.Errorf("expected '%s', got: %v", expectedErr, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStopShellStream_NonExistent(t *testing.T) {
|
||||
log := newTestLogger()
|
||||
handlers := NewShellHandlers(nil, log)
|
||||
|
||||
// StopShellStream should not panic when called with non-existent sessionID
|
||||
handlers.StopShellStream("non-existent-session")
|
||||
// No panic = success
|
||||
}
|
||||
|
||||
func TestNewShellHandlers_WithManager(t *testing.T) {
|
||||
log := newTestLogger()
|
||||
mgr := newTestManager()
|
||||
|
||||
@@ -108,3 +108,32 @@ type PermissionRequestEventPayload struct {
|
||||
ActionDetails map[string]interface{} `json:"action_details,omitempty"`
|
||||
}
|
||||
|
||||
// ShellOutputEventPayload is the payload for shell output events.
|
||||
type ShellOutputEventPayload struct {
|
||||
TaskID string `json:"task_id"`
|
||||
SessionID string `json:"session_id"`
|
||||
AgentID string `json:"agent_id"`
|
||||
Type string `json:"type"` // Always "output" for shell output events
|
||||
Data string `json:"data"`
|
||||
Timestamp string `json:"timestamp"`
|
||||
}
|
||||
|
||||
// GetSessionID returns the session ID for this event (used by event routing).
|
||||
func (p ShellOutputEventPayload) GetSessionID() string {
|
||||
return p.SessionID
|
||||
}
|
||||
|
||||
// ShellExitEventPayload is the payload for shell exit events.
|
||||
type ShellExitEventPayload struct {
|
||||
TaskID string `json:"task_id"`
|
||||
SessionID string `json:"session_id"`
|
||||
AgentID string `json:"agent_id"`
|
||||
Type string `json:"type"` // Always "exit" for shell exit events
|
||||
Code int `json:"code"` // Exit code
|
||||
Timestamp string `json:"timestamp"`
|
||||
}
|
||||
|
||||
// GetSessionID returns the session ID for this event (used by event routing).
|
||||
func (p ShellExitEventPayload) GetSessionID() string {
|
||||
return p.SessionID
|
||||
}
|
||||
|
||||
@@ -265,3 +265,61 @@ func (p *EventPublisher) PublishPermissionRequest(execution *AgentExecution, eve
|
||||
zap.String("title", event.PermissionTitle))
|
||||
}
|
||||
}
|
||||
|
||||
// PublishShellOutput publishes a shell output event to the event bus.
|
||||
func (p *EventPublisher) PublishShellOutput(execution *AgentExecution, data string) {
|
||||
if p.eventBus == nil {
|
||||
return
|
||||
}
|
||||
|
||||
sessionID := execution.SessionID
|
||||
|
||||
payload := ShellOutputEventPayload{
|
||||
TaskID: execution.TaskID,
|
||||
SessionID: sessionID,
|
||||
AgentID: execution.ID,
|
||||
Type: "output",
|
||||
Data: data,
|
||||
Timestamp: time.Now().UTC().Format(time.RFC3339Nano),
|
||||
}
|
||||
|
||||
event := bus.NewEvent(events.ShellOutput, "agent-manager", payload)
|
||||
subject := events.BuildShellOutputSubject(sessionID)
|
||||
|
||||
if err := p.eventBus.Publish(context.Background(), subject, event); err != nil {
|
||||
p.logger.Error("failed to publish shell output event",
|
||||
zap.String("instance_id", execution.ID),
|
||||
zap.String("task_id", execution.TaskID),
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
// PublishShellExit publishes a shell exit event to the event bus.
|
||||
func (p *EventPublisher) PublishShellExit(execution *AgentExecution, exitCode int) {
|
||||
if p.eventBus == nil {
|
||||
return
|
||||
}
|
||||
|
||||
sessionID := execution.SessionID
|
||||
|
||||
payload := ShellExitEventPayload{
|
||||
TaskID: execution.TaskID,
|
||||
SessionID: sessionID,
|
||||
AgentID: execution.ID,
|
||||
Type: "exit",
|
||||
Code: exitCode,
|
||||
Timestamp: time.Now().UTC().Format(time.RFC3339Nano),
|
||||
}
|
||||
|
||||
event := bus.NewEvent(events.ShellExit, "agent-manager", payload)
|
||||
subject := events.BuildShellExitSubject(sessionID)
|
||||
|
||||
if err := p.eventBus.Publish(context.Background(), subject, event); err != nil {
|
||||
p.logger.Error("failed to publish shell exit event",
|
||||
zap.String("instance_id", execution.ID),
|
||||
zap.String("task_id", execution.TaskID),
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -44,9 +44,6 @@ type Manager struct {
|
||||
eventPublisher *EventPublisher // Publishes lifecycle events
|
||||
containerManager *ContainerManager // Manages Docker containers (optional, nil for non-Docker runtimes)
|
||||
|
||||
// Shell stream starter for auto-starting shell streams
|
||||
shellStreamStarter ShellStreamStarter
|
||||
|
||||
// Workspace info provider for on-demand instance creation
|
||||
workspaceInfoProvider WorkspaceInfoProvider
|
||||
|
||||
@@ -100,9 +97,11 @@ func NewManager(
|
||||
|
||||
// Initialize stream manager with callbacks that delegate to manager methods
|
||||
mgr.streamManager = NewStreamManager(log, StreamCallbacks{
|
||||
OnAgentEvent: mgr.handleAgentEvent,
|
||||
OnGitStatus: mgr.handleGitStatusUpdate,
|
||||
OnFileChange: mgr.handleFileChangeNotification,
|
||||
OnAgentEvent: mgr.handleAgentEvent,
|
||||
OnGitStatus: mgr.handleGitStatusUpdate,
|
||||
OnFileChange: mgr.handleFileChangeNotification,
|
||||
OnShellOutput: mgr.handleShellOutput,
|
||||
OnShellExit: mgr.handleShellExit,
|
||||
})
|
||||
|
||||
// Set session manager dependencies for full orchestration
|
||||
@@ -120,11 +119,6 @@ func (m *Manager) SetWorktreeManager(worktreeMgr *worktree.Manager) {
|
||||
m.worktreeMgr = worktreeMgr
|
||||
}
|
||||
|
||||
// SetShellStreamStarter sets the shell stream starter
|
||||
func (m *Manager) SetShellStreamStarter(starter ShellStreamStarter) {
|
||||
m.shellStreamStarter = starter
|
||||
}
|
||||
|
||||
// SetWorkspaceInfoProvider sets the provider for workspace information
|
||||
func (m *Manager) SetWorkspaceInfoProvider(provider WorkspaceInfoProvider) {
|
||||
m.workspaceInfoProvider = provider
|
||||
@@ -856,6 +850,16 @@ func (m *Manager) handleFileChangeNotification(execution *AgentExecution, notifi
|
||||
m.eventPublisher.PublishFileChange(execution, notification)
|
||||
}
|
||||
|
||||
// handleShellOutput processes shell output from the workspace stream
|
||||
func (m *Manager) handleShellOutput(execution *AgentExecution, data string) {
|
||||
m.eventPublisher.PublishShellOutput(execution, data)
|
||||
}
|
||||
|
||||
// handleShellExit processes shell exit events from the workspace stream
|
||||
func (m *Manager) handleShellExit(execution *AgentExecution, code int) {
|
||||
m.eventPublisher.PublishShellExit(execution, code)
|
||||
}
|
||||
|
||||
// flushMessageBuffer extracts any accumulated message from the buffer and returns it.
|
||||
// This is called when a tool use starts or on complete to get the agent's response.
|
||||
func (m *Manager) flushMessageBuffer(execution *AgentExecution) string {
|
||||
|
||||
@@ -15,6 +15,8 @@ type StreamCallbacks struct {
|
||||
OnAgentEvent func(execution *AgentExecution, event agentctl.AgentEvent)
|
||||
OnGitStatus func(execution *AgentExecution, update *agentctl.GitStatusUpdate)
|
||||
OnFileChange func(execution *AgentExecution, notification *agentctl.FileChangeNotification)
|
||||
OnShellOutput func(execution *AgentExecution, data string)
|
||||
OnShellExit func(execution *AgentExecution, code int)
|
||||
}
|
||||
|
||||
// StreamManager manages WebSocket streams to agent executions
|
||||
@@ -35,8 +37,7 @@ func NewStreamManager(log *logger.Logger, callbacks StreamCallbacks) *StreamMana
|
||||
// If ready is non-nil, it will be closed when the updates stream is connected.
|
||||
func (sm *StreamManager) ConnectAll(execution *AgentExecution, ready chan<- struct{}) {
|
||||
go sm.connectUpdatesStream(execution, ready)
|
||||
go sm.connectGitStatusStream(execution)
|
||||
go sm.connectFileChangesStream(execution)
|
||||
go sm.connectWorkspaceStream(execution)
|
||||
}
|
||||
|
||||
// ReconnectAll reconnects to all streams (used after backend restart).
|
||||
@@ -92,10 +93,8 @@ func (sm *StreamManager) connectUpdatesStream(execution *AgentExecution, ready c
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
// connectGitStatusStream handles git status stream with retry logic
|
||||
func (sm *StreamManager) connectGitStatusStream(execution *AgentExecution) {
|
||||
// connectWorkspaceStream handles the unified workspace stream with retry logic
|
||||
func (sm *StreamManager) connectWorkspaceStream(execution *AgentExecution) {
|
||||
ctx := context.Background()
|
||||
|
||||
// Retry connection with exponential backoff
|
||||
@@ -103,67 +102,64 @@ func (sm *StreamManager) connectGitStatusStream(execution *AgentExecution) {
|
||||
backoff := 1 * time.Second
|
||||
|
||||
for attempt := 1; attempt <= maxRetries; attempt++ {
|
||||
err := execution.agentctl.StreamGitStatus(ctx, func(update *agentctl.GitStatusUpdate) {
|
||||
if sm.callbacks.OnGitStatus != nil {
|
||||
sm.callbacks.OnGitStatus(execution, update)
|
||||
callbacks := agentctl.WorkspaceStreamCallbacks{
|
||||
OnShellOutput: func(data string) {
|
||||
if sm.callbacks.OnShellOutput != nil {
|
||||
sm.callbacks.OnShellOutput(execution, data)
|
||||
}
|
||||
},
|
||||
OnShellExit: func(code int) {
|
||||
if sm.callbacks.OnShellExit != nil {
|
||||
sm.callbacks.OnShellExit(execution, code)
|
||||
}
|
||||
},
|
||||
OnGitStatus: func(update *agentctl.GitStatusUpdate) {
|
||||
if sm.callbacks.OnGitStatus != nil {
|
||||
sm.callbacks.OnGitStatus(execution, update)
|
||||
}
|
||||
},
|
||||
OnFileChange: func(notification *agentctl.FileChangeNotification) {
|
||||
if sm.callbacks.OnFileChange != nil {
|
||||
sm.callbacks.OnFileChange(execution, notification)
|
||||
}
|
||||
},
|
||||
OnConnected: func() {
|
||||
sm.logger.Debug("workspace stream connected",
|
||||
zap.String("instance_id", execution.ID))
|
||||
},
|
||||
OnError: func(err string) {
|
||||
sm.logger.Debug("workspace stream error",
|
||||
zap.String("instance_id", execution.ID),
|
||||
zap.String("error", err))
|
||||
},
|
||||
}
|
||||
|
||||
ws, err := execution.agentctl.StreamWorkspace(ctx, callbacks)
|
||||
if err != nil {
|
||||
sm.logger.Debug("workspace stream connection failed, retrying",
|
||||
zap.String("instance_id", execution.ID),
|
||||
zap.Int("attempt", attempt),
|
||||
zap.Int("max_retries", maxRetries),
|
||||
zap.Error(err))
|
||||
|
||||
if attempt < maxRetries {
|
||||
time.Sleep(backoff)
|
||||
backoff *= 2 // Exponential backoff
|
||||
}
|
||||
})
|
||||
|
||||
if err == nil {
|
||||
// Connection closed normally
|
||||
return
|
||||
continue
|
||||
}
|
||||
|
||||
sm.logger.Debug("git status stream connection failed, retrying",
|
||||
zap.String("instance_id", execution.ID),
|
||||
zap.Int("attempt", attempt),
|
||||
zap.Int("max_retries", maxRetries),
|
||||
zap.Error(err))
|
||||
// Store the workspace stream on the execution for shell I/O
|
||||
execution.SetWorkspaceStream(ws)
|
||||
sm.logger.Info("connected to unified workspace stream",
|
||||
zap.String("instance_id", execution.ID))
|
||||
|
||||
if attempt < maxRetries {
|
||||
time.Sleep(backoff)
|
||||
backoff *= 2 // Exponential backoff
|
||||
}
|
||||
// Wait for the stream to close (it stays open until disconnected)
|
||||
<-ws.Done()
|
||||
return
|
||||
}
|
||||
|
||||
sm.logger.Error("failed to connect to git status stream after retries",
|
||||
zap.String("instance_id", execution.ID),
|
||||
zap.Int("max_retries", maxRetries))
|
||||
}
|
||||
|
||||
// connectFileChangesStream handles file changes stream with retry logic
|
||||
func (sm *StreamManager) connectFileChangesStream(execution *AgentExecution) {
|
||||
ctx := context.Background()
|
||||
|
||||
// Retry connection with exponential backoff
|
||||
maxRetries := 5
|
||||
backoff := 1 * time.Second
|
||||
|
||||
for attempt := 1; attempt <= maxRetries; attempt++ {
|
||||
err := execution.agentctl.StreamFileChanges(ctx, func(notification *agentctl.FileChangeNotification) {
|
||||
if sm.callbacks.OnFileChange != nil {
|
||||
sm.callbacks.OnFileChange(execution, notification)
|
||||
}
|
||||
})
|
||||
|
||||
if err == nil {
|
||||
// Connection closed normally
|
||||
return
|
||||
}
|
||||
|
||||
sm.logger.Debug("file changes stream connection failed, retrying",
|
||||
zap.String("instance_id", execution.ID),
|
||||
zap.Int("attempt", attempt),
|
||||
zap.Int("max_retries", maxRetries),
|
||||
zap.Error(err))
|
||||
|
||||
if attempt < maxRetries {
|
||||
time.Sleep(backoff)
|
||||
backoff *= 2 // Exponential backoff
|
||||
}
|
||||
}
|
||||
|
||||
sm.logger.Error("failed to connect to file changes stream after retries",
|
||||
sm.logger.Error("failed to connect to workspace stream after retries",
|
||||
zap.String("instance_id", execution.ID),
|
||||
zap.Int("max_retries", maxRetries))
|
||||
}
|
||||
|
||||
@@ -37,6 +37,10 @@ type AgentExecution struct {
|
||||
// agentctl client for this execution
|
||||
agentctl *agentctl.Client
|
||||
|
||||
// Unified workspace stream for shell I/O, git status, and file changes
|
||||
workspaceStream *agentctl.WorkspaceStream
|
||||
workspaceStreamMu sync.RWMutex
|
||||
|
||||
// Standalone mode info (when not using Docker)
|
||||
standaloneInstanceID string // Instance ID in standalone agentctl
|
||||
standalonePort int // Port of the standalone execution
|
||||
@@ -53,6 +57,20 @@ func (ae *AgentExecution) GetAgentCtlClient() *agentctl.Client {
|
||||
return ae.agentctl
|
||||
}
|
||||
|
||||
// SetWorkspaceStream sets the unified workspace stream for this execution
|
||||
func (ae *AgentExecution) SetWorkspaceStream(ws *agentctl.WorkspaceStream) {
|
||||
ae.workspaceStreamMu.Lock()
|
||||
defer ae.workspaceStreamMu.Unlock()
|
||||
ae.workspaceStream = ws
|
||||
}
|
||||
|
||||
// GetWorkspaceStream returns the unified workspace stream for this execution
|
||||
func (ae *AgentExecution) GetWorkspaceStream() *agentctl.WorkspaceStream {
|
||||
ae.workspaceStreamMu.RLock()
|
||||
defer ae.workspaceStreamMu.RUnlock()
|
||||
return ae.workspaceStream
|
||||
}
|
||||
|
||||
// LaunchRequest contains parameters for launching an agent
|
||||
type LaunchRequest struct {
|
||||
TaskID string
|
||||
@@ -94,11 +112,6 @@ type ProfileResolver interface {
|
||||
ResolveProfile(ctx context.Context, profileID string) (*AgentProfileInfo, error)
|
||||
}
|
||||
|
||||
// ShellStreamStarter is called to start shell streaming for an agent execution
|
||||
type ShellStreamStarter interface {
|
||||
StartShellStream(ctx context.Context, sessionID string) error
|
||||
}
|
||||
|
||||
// WorkspaceInfo contains information about a task's workspace for on-demand execution creation
|
||||
type WorkspaceInfo struct {
|
||||
TaskID string
|
||||
|
||||
@@ -23,11 +23,9 @@ type Client struct {
|
||||
logger *logger.Logger
|
||||
|
||||
// WebSocket connections for streaming
|
||||
acpConn *websocket.Conn
|
||||
gitStatusConn *websocket.Conn
|
||||
fileChangesConn *websocket.Conn
|
||||
shellConn *websocket.Conn
|
||||
mu sync.RWMutex
|
||||
acpConn *websocket.Conn
|
||||
workspaceStreamConn *websocket.Conn
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// StatusResponse from agentctl
|
||||
@@ -268,119 +266,6 @@ type (
|
||||
ShellBufferResponse = types.ShellBufferResponse
|
||||
)
|
||||
|
||||
// StreamGitStatus opens a WebSocket connection for streaming git status updates
|
||||
func (c *Client) StreamGitStatus(ctx context.Context, handler func(*GitStatusUpdate)) error {
|
||||
wsURL := "ws" + c.baseURL[4:] + "/api/v1/workspace/git-status/stream"
|
||||
|
||||
conn, _, err := websocket.DefaultDialer.DialContext(ctx, wsURL, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to connect to git status stream: %w", err)
|
||||
}
|
||||
|
||||
c.mu.Lock()
|
||||
c.gitStatusConn = conn
|
||||
c.mu.Unlock()
|
||||
|
||||
c.logger.Info("connected to git status stream", zap.String("url", wsURL))
|
||||
|
||||
// Read messages in a goroutine
|
||||
go func() {
|
||||
defer func() {
|
||||
c.mu.Lock()
|
||||
c.gitStatusConn = nil
|
||||
c.mu.Unlock()
|
||||
if err := conn.Close(); err != nil {
|
||||
c.logger.Debug("failed to close git status websocket", zap.Error(err))
|
||||
}
|
||||
}()
|
||||
|
||||
for {
|
||||
_, message, err := conn.ReadMessage()
|
||||
if err != nil {
|
||||
if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
|
||||
c.logger.Info("git status stream closed normally")
|
||||
} else {
|
||||
c.logger.Debug("git status stream error", zap.Error(err))
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
var update GitStatusUpdate
|
||||
if err := json.Unmarshal(message, &update); err != nil {
|
||||
c.logger.Warn("failed to parse git status update", zap.Error(err))
|
||||
continue
|
||||
}
|
||||
|
||||
handler(&update)
|
||||
}
|
||||
}()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// CloseGitStatusStream closes the git status stream connection
|
||||
func (c *Client) CloseGitStatusStream() {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
if c.gitStatusConn != nil {
|
||||
if err := c.gitStatusConn.Close(); err != nil {
|
||||
c.logger.Debug("failed to close git status stream", zap.Error(err))
|
||||
}
|
||||
c.gitStatusConn = nil
|
||||
}
|
||||
}
|
||||
|
||||
// StreamFileChanges opens a WebSocket connection for streaming file change notifications
|
||||
func (c *Client) StreamFileChanges(ctx context.Context, handler func(*FileChangeNotification)) error {
|
||||
wsURL := "ws" + c.baseURL[4:] + "/api/v1/workspace/file-changes/stream"
|
||||
|
||||
conn, _, err := websocket.DefaultDialer.DialContext(ctx, wsURL, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to connect to file changes stream: %w", err)
|
||||
}
|
||||
|
||||
c.mu.Lock()
|
||||
c.fileChangesConn = conn
|
||||
c.mu.Unlock()
|
||||
|
||||
c.logger.Info("connected to file changes stream", zap.String("url", wsURL))
|
||||
|
||||
// Read messages in a goroutine
|
||||
go func() {
|
||||
defer func() {
|
||||
c.mu.Lock()
|
||||
c.fileChangesConn = nil
|
||||
c.mu.Unlock()
|
||||
if err := conn.Close(); err != nil {
|
||||
c.logger.Debug("failed to close file changes websocket", zap.Error(err))
|
||||
}
|
||||
}()
|
||||
|
||||
for {
|
||||
_, message, err := conn.ReadMessage()
|
||||
if err != nil {
|
||||
if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
|
||||
c.logger.Info("file changes stream closed normally")
|
||||
} else {
|
||||
c.logger.Debug("file changes stream error", zap.Error(err))
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
var notification FileChangeNotification
|
||||
if err := json.Unmarshal(message, ¬ification); err != nil {
|
||||
c.logger.Warn("failed to parse file change notification", zap.Error(err))
|
||||
continue
|
||||
}
|
||||
|
||||
handler(¬ification)
|
||||
}
|
||||
}()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// RequestFileTree requests a file tree via HTTP GET
|
||||
func (c *Client) RequestFileTree(ctx context.Context, path string, depth int) (*FileTreeResponse, error) {
|
||||
url := fmt.Sprintf("%s/api/v1/workspace/tree?path=%s&depth=%d", c.baseURL, path, depth)
|
||||
@@ -443,25 +328,10 @@ func (c *Client) RequestFileContent(ctx context.Context, path string) (*FileCont
|
||||
return &response, nil
|
||||
}
|
||||
|
||||
// CloseFileChangesStream closes the file changes stream connection
|
||||
func (c *Client) CloseFileChangesStream() {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
if c.fileChangesConn != nil {
|
||||
if err := c.fileChangesConn.Close(); err != nil {
|
||||
c.logger.Debug("failed to close file changes stream", zap.Error(err))
|
||||
}
|
||||
c.fileChangesConn = nil
|
||||
}
|
||||
}
|
||||
|
||||
// Close closes all connections
|
||||
func (c *Client) Close() {
|
||||
c.CloseUpdatesStream()
|
||||
c.CloseGitStatusStream()
|
||||
c.CloseFileChangesStream()
|
||||
c.CloseShellStream()
|
||||
c.CloseWorkspaceStream()
|
||||
}
|
||||
|
||||
// Note: ShellStatusResponse, ShellMessage, ShellBufferResponse are re-exported
|
||||
@@ -517,67 +387,186 @@ func (c *Client) ShellBuffer(ctx context.Context) (string, error) {
|
||||
return result.Data, nil
|
||||
}
|
||||
|
||||
// StreamShell connects to the shell WebSocket stream and returns channels for I/O
|
||||
func (c *Client) StreamShell(ctx context.Context) (<-chan ShellMessage, chan<- ShellMessage, error) {
|
||||
// WorkspaceStreamCallbacks defines callbacks for workspace stream events
|
||||
type WorkspaceStreamCallbacks struct {
|
||||
OnShellOutput func(data string)
|
||||
OnShellExit func(code int)
|
||||
OnGitStatus func(update *GitStatusUpdate)
|
||||
OnFileChange func(notification *FileChangeNotification)
|
||||
OnFileList func(update *FileListUpdate)
|
||||
OnConnected func()
|
||||
OnError func(err string)
|
||||
}
|
||||
|
||||
// WorkspaceStream represents an active workspace stream connection
|
||||
type WorkspaceStream struct {
|
||||
conn *websocket.Conn
|
||||
inputCh chan types.WorkspaceStreamMessage
|
||||
closeCh chan struct{}
|
||||
closeOnce sync.Once
|
||||
logger *logger.Logger
|
||||
}
|
||||
|
||||
// StreamWorkspace opens a unified WebSocket connection for all workspace events
|
||||
func (c *Client) StreamWorkspace(ctx context.Context, callbacks WorkspaceStreamCallbacks) (*WorkspaceStream, error) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
if c.shellConn != nil {
|
||||
return nil, nil, fmt.Errorf("shell stream already connected")
|
||||
if c.workspaceStreamConn != nil {
|
||||
c.mu.Unlock()
|
||||
return nil, fmt.Errorf("workspace stream already connected")
|
||||
}
|
||||
c.mu.Unlock()
|
||||
|
||||
wsURL := "ws" + c.baseURL[4:] + "/api/v1/shell/stream"
|
||||
wsURL := "ws" + c.baseURL[4:] + "/api/v1/workspace/stream"
|
||||
conn, _, err := websocket.DefaultDialer.DialContext(ctx, wsURL, nil)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("failed to connect shell stream: %w", err)
|
||||
return nil, fmt.Errorf("failed to connect to workspace stream: %w", err)
|
||||
}
|
||||
c.shellConn = conn
|
||||
|
||||
outputCh := make(chan ShellMessage, 256)
|
||||
inputCh := make(chan ShellMessage, 64)
|
||||
c.mu.Lock()
|
||||
c.workspaceStreamConn = conn
|
||||
c.mu.Unlock()
|
||||
|
||||
// Read from WebSocket
|
||||
c.logger.Info("connected to workspace stream", zap.String("url", wsURL))
|
||||
|
||||
ws := &WorkspaceStream{
|
||||
conn: conn,
|
||||
inputCh: make(chan types.WorkspaceStreamMessage, 64),
|
||||
closeCh: make(chan struct{}),
|
||||
logger: c.logger,
|
||||
}
|
||||
|
||||
// Read goroutine - dispatches to callbacks
|
||||
go func() {
|
||||
defer close(outputCh)
|
||||
defer func() {
|
||||
c.mu.Lock()
|
||||
c.workspaceStreamConn = nil
|
||||
c.mu.Unlock()
|
||||
ws.Close()
|
||||
}()
|
||||
|
||||
for {
|
||||
var msg ShellMessage
|
||||
var msg types.WorkspaceStreamMessage
|
||||
if err := conn.ReadJSON(&msg); err != nil {
|
||||
if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
|
||||
c.logger.Debug("shell stream read error", zap.Error(err))
|
||||
c.logger.Debug("workspace stream read error", zap.Error(err))
|
||||
}
|
||||
return
|
||||
}
|
||||
select {
|
||||
case outputCh <- msg:
|
||||
default:
|
||||
// Channel full, skip
|
||||
|
||||
switch msg.Type {
|
||||
case types.WorkspaceMessageTypeShellOutput:
|
||||
if callbacks.OnShellOutput != nil {
|
||||
callbacks.OnShellOutput(msg.Data)
|
||||
}
|
||||
case types.WorkspaceMessageTypeShellExit:
|
||||
if callbacks.OnShellExit != nil {
|
||||
callbacks.OnShellExit(msg.Code)
|
||||
}
|
||||
case types.WorkspaceMessageTypeGitStatus:
|
||||
if callbacks.OnGitStatus != nil && msg.GitStatus != nil {
|
||||
callbacks.OnGitStatus(msg.GitStatus)
|
||||
}
|
||||
case types.WorkspaceMessageTypeFileChange:
|
||||
if callbacks.OnFileChange != nil && msg.FileChange != nil {
|
||||
callbacks.OnFileChange(msg.FileChange)
|
||||
}
|
||||
case types.WorkspaceMessageTypeFileList:
|
||||
if callbacks.OnFileList != nil && msg.FileList != nil {
|
||||
callbacks.OnFileList(msg.FileList)
|
||||
}
|
||||
case types.WorkspaceMessageTypeConnected:
|
||||
if callbacks.OnConnected != nil {
|
||||
callbacks.OnConnected()
|
||||
}
|
||||
case types.WorkspaceMessageTypeError:
|
||||
if callbacks.OnError != nil {
|
||||
callbacks.OnError(msg.Error)
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// Write to WebSocket
|
||||
// Write goroutine - sends from inputCh
|
||||
go func() {
|
||||
for msg := range inputCh {
|
||||
if err := conn.WriteJSON(msg); err != nil {
|
||||
c.logger.Debug("shell stream write error", zap.Error(err))
|
||||
for {
|
||||
select {
|
||||
case <-ws.closeCh:
|
||||
return
|
||||
case msg, ok := <-ws.inputCh:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if err := conn.WriteJSON(msg); err != nil {
|
||||
ws.logger.Debug("workspace stream write error", zap.Error(err))
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
return outputCh, inputCh, nil
|
||||
return ws, nil
|
||||
}
|
||||
|
||||
// CloseShellStream closes the shell stream connection
|
||||
func (c *Client) CloseShellStream() {
|
||||
// WriteShellInput sends input to the shell through the workspace stream
|
||||
func (ws *WorkspaceStream) WriteShellInput(data string) error {
|
||||
msg := types.NewWorkspaceShellInput(data)
|
||||
select {
|
||||
case ws.inputCh <- msg:
|
||||
return nil
|
||||
case <-ws.closeCh:
|
||||
return fmt.Errorf("workspace stream closed")
|
||||
}
|
||||
}
|
||||
|
||||
// ResizeShell sends a shell resize command through the workspace stream
|
||||
func (ws *WorkspaceStream) ResizeShell(cols, rows int) error {
|
||||
msg := types.NewWorkspaceShellResize(cols, rows)
|
||||
select {
|
||||
case ws.inputCh <- msg:
|
||||
return nil
|
||||
case <-ws.closeCh:
|
||||
return fmt.Errorf("workspace stream closed")
|
||||
}
|
||||
}
|
||||
|
||||
// Ping sends a ping message through the workspace stream
|
||||
func (ws *WorkspaceStream) Ping() error {
|
||||
msg := types.NewWorkspacePing()
|
||||
select {
|
||||
case ws.inputCh <- msg:
|
||||
return nil
|
||||
case <-ws.closeCh:
|
||||
return fmt.Errorf("workspace stream closed")
|
||||
}
|
||||
}
|
||||
|
||||
// Close closes the workspace stream
|
||||
func (ws *WorkspaceStream) Close() {
|
||||
ws.closeOnce.Do(func() {
|
||||
close(ws.closeCh)
|
||||
if ws.conn != nil {
|
||||
if err := ws.conn.Close(); err != nil {
|
||||
ws.logger.Debug("failed to close workspace stream connection", zap.Error(err))
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// Done returns a channel that is closed when the stream is closed
|
||||
func (ws *WorkspaceStream) Done() <-chan struct{} {
|
||||
return ws.closeCh
|
||||
}
|
||||
|
||||
// CloseWorkspaceStream closes the workspace stream connection
|
||||
func (c *Client) CloseWorkspaceStream() {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
if c.shellConn != nil {
|
||||
if err := c.shellConn.Close(); err != nil {
|
||||
c.logger.Debug("failed to close shell stream", zap.Error(err))
|
||||
if c.workspaceStreamConn != nil {
|
||||
if err := c.workspaceStreamConn.Close(); err != nil {
|
||||
c.logger.Debug("failed to close workspace stream", zap.Error(err))
|
||||
}
|
||||
c.shellConn = nil
|
||||
c.workspaceStreamConn = nil
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -74,18 +74,15 @@ func (s *Server) setupRoutes() {
|
||||
// Permission request handling
|
||||
api.POST("/acp/permissions/respond", s.handlePermissionRespond)
|
||||
|
||||
// Workspace monitoring (git status, files)
|
||||
api.GET("/workspace/git-status/stream", s.handleGitStatusStreamWS)
|
||||
api.GET("/workspace/files/stream", s.handleFilesStreamWS)
|
||||
api.GET("/workspace/file-changes/stream", s.handleFileChangesStreamWS)
|
||||
// Unified workspace stream (git status, files, shell)
|
||||
api.GET("/workspace/stream", s.handleWorkspaceStreamWS)
|
||||
|
||||
// Workspace file operations (simple HTTP)
|
||||
api.GET("/workspace/tree", s.handleFileTree)
|
||||
api.GET("/workspace/file/content", s.handleFileContent)
|
||||
|
||||
// Shell access
|
||||
// Shell access (HTTP endpoints only - streaming is via /workspace/stream)
|
||||
api.GET("/shell/status", s.handleShellStatus)
|
||||
api.GET("/shell/stream", s.handleShellStreamWS)
|
||||
api.GET("/shell/buffer", s.handleShellBuffer)
|
||||
}
|
||||
}
|
||||
@@ -290,74 +287,3 @@ func (s *Server) handleShellBuffer(c *gin.Context) {
|
||||
Data: string(buffered),
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) handleShellStreamWS(c *gin.Context) {
|
||||
shell := s.procMgr.Shell()
|
||||
if shell == nil {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "shell not available"})
|
||||
return
|
||||
}
|
||||
|
||||
conn, err := s.upgrader.Upgrade(c.Writer, c.Request, nil)
|
||||
if err != nil {
|
||||
s.logger.Error("WebSocket upgrade failed", zap.Error(err))
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
if err := conn.Close(); err != nil {
|
||||
s.logger.Debug("failed to close shell websocket", zap.Error(err))
|
||||
}
|
||||
}()
|
||||
|
||||
s.logger.Info("Shell stream WebSocket connected")
|
||||
|
||||
// Create output channel and subscribe to shell output
|
||||
outputCh := make(chan []byte, 256)
|
||||
shell.Subscribe(outputCh)
|
||||
defer shell.Unsubscribe(outputCh)
|
||||
|
||||
// Note: Buffered output is sent via GET /api/v1/shell/buffer when clients subscribe
|
||||
// This avoids duplicate output when clients reconnect
|
||||
|
||||
// Done channel to signal goroutine shutdown
|
||||
done := make(chan struct{})
|
||||
defer close(done)
|
||||
|
||||
// Read input from WebSocket and write to shell
|
||||
go func() {
|
||||
for {
|
||||
select {
|
||||
case <-done:
|
||||
return
|
||||
default:
|
||||
var msg ShellMessage
|
||||
if err := conn.ReadJSON(&msg); err != nil {
|
||||
s.logger.Debug("shell WebSocket read error", zap.Error(err))
|
||||
return
|
||||
}
|
||||
|
||||
switch msg.Type {
|
||||
case "input":
|
||||
if _, err := shell.Write([]byte(msg.Data)); err != nil {
|
||||
s.logger.Debug("shell write error", zap.Error(err))
|
||||
}
|
||||
case "ping":
|
||||
// Respond with pong
|
||||
if err := conn.WriteJSON(ShellMessage{Type: "pong"}); err != nil {
|
||||
s.logger.Debug("shell pong write error", zap.Error(err))
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// Write output from shell to WebSocket
|
||||
for data := range outputCh {
|
||||
msg := ShellMessage{Type: "output", Data: string(data)}
|
||||
if err := conn.WriteJSON(msg); err != nil {
|
||||
s.logger.Debug("shell WebSocket write error", zap.Error(err))
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4,13 +4,13 @@ import (
|
||||
"encoding/json"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
"github.com/kandev/kandev/internal/agentctl/types"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// handleGitStatusStreamWS streams git status updates via WebSocket
|
||||
func (s *Server) handleGitStatusStreamWS(c *gin.Context) {
|
||||
// handleWorkspaceStreamWS handles the unified workspace stream WebSocket endpoint.
|
||||
// It streams git status, file changes, file lists, and shell I/O over a single connection.
|
||||
func (s *Server) handleWorkspaceStreamWS(c *gin.Context) {
|
||||
conn, err := s.upgrader.Upgrade(c.Writer, c.Request, nil)
|
||||
if err != nil {
|
||||
s.logger.Error("WebSocket upgrade failed", zap.Error(err))
|
||||
@@ -18,143 +18,99 @@ func (s *Server) handleGitStatusStreamWS(c *gin.Context) {
|
||||
}
|
||||
defer func() {
|
||||
if err := conn.Close(); err != nil {
|
||||
s.logger.Debug("failed to close git status websocket", zap.Error(err))
|
||||
s.logger.Debug("failed to close workspace stream websocket", zap.Error(err))
|
||||
}
|
||||
}()
|
||||
|
||||
// Subscribe to git status updates
|
||||
sub := s.procMgr.GetWorkspaceTracker().SubscribeGitStatus()
|
||||
defer s.procMgr.GetWorkspaceTracker().UnsubscribeGitStatus(sub)
|
||||
s.logger.Info("Workspace stream WebSocket connected")
|
||||
|
||||
// Handle WebSocket close
|
||||
closeCh := make(chan struct{})
|
||||
go func() {
|
||||
for {
|
||||
_, _, err := conn.ReadMessage()
|
||||
if err != nil {
|
||||
close(closeCh)
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
// Subscribe to unified workspace updates
|
||||
sub := s.procMgr.GetWorkspaceTracker().SubscribeWorkspaceStream()
|
||||
defer s.procMgr.GetWorkspaceTracker().UnsubscribeWorkspaceStream(sub)
|
||||
|
||||
// Stream git status updates to client
|
||||
for {
|
||||
select {
|
||||
case update, ok := <-sub:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
data, err := json.Marshal(update)
|
||||
if err != nil {
|
||||
s.logger.Error("failed to marshal git status update", zap.Error(err))
|
||||
continue
|
||||
}
|
||||
if err := conn.WriteMessage(websocket.TextMessage, data); err != nil {
|
||||
return
|
||||
}
|
||||
case <-closeCh:
|
||||
return
|
||||
}
|
||||
// Get shell for input handling (may be nil)
|
||||
shell := s.procMgr.Shell()
|
||||
|
||||
// Subscribe to shell output if shell is available
|
||||
var shellOutputCh chan []byte
|
||||
if shell != nil {
|
||||
shellOutputCh = make(chan []byte, 256)
|
||||
shell.Subscribe(shellOutputCh)
|
||||
defer shell.Unsubscribe(shellOutputCh)
|
||||
}
|
||||
}
|
||||
|
||||
// handleFilesStreamWS streams file listing updates via WebSocket
|
||||
func (s *Server) handleFilesStreamWS(c *gin.Context) {
|
||||
conn, err := s.upgrader.Upgrade(c.Writer, c.Request, nil)
|
||||
if err != nil {
|
||||
s.logger.Error("WebSocket upgrade failed", zap.Error(err))
|
||||
// Send connected message
|
||||
connectedMsg := types.NewWorkspaceConnected()
|
||||
if err := conn.WriteJSON(connectedMsg); err != nil {
|
||||
s.logger.Debug("failed to send connected message", zap.Error(err))
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
if err := conn.Close(); err != nil {
|
||||
s.logger.Debug("failed to close files websocket", zap.Error(err))
|
||||
}
|
||||
}()
|
||||
|
||||
// Subscribe to file updates
|
||||
sub := s.procMgr.GetWorkspaceTracker().SubscribeFiles()
|
||||
defer s.procMgr.GetWorkspaceTracker().UnsubscribeFiles(sub)
|
||||
// Done channel to signal goroutine shutdown
|
||||
done := make(chan struct{})
|
||||
defer close(done)
|
||||
|
||||
// Handle WebSocket close
|
||||
closeCh := make(chan struct{})
|
||||
// Handle incoming messages (shell_input, shell_resize, ping) in a goroutine
|
||||
go func() {
|
||||
for {
|
||||
_, _, err := conn.ReadMessage()
|
||||
if err != nil {
|
||||
close(closeCh)
|
||||
select {
|
||||
case <-done:
|
||||
return
|
||||
default:
|
||||
var msg types.WorkspaceStreamMessage
|
||||
if err := conn.ReadJSON(&msg); err != nil {
|
||||
s.logger.Debug("workspace stream WebSocket read error", zap.Error(err))
|
||||
return
|
||||
}
|
||||
|
||||
switch msg.Type {
|
||||
case types.WorkspaceMessageTypeShellInput:
|
||||
if shell != nil {
|
||||
if _, err := shell.Write([]byte(msg.Data)); err != nil {
|
||||
s.logger.Debug("shell write error", zap.Error(err))
|
||||
}
|
||||
}
|
||||
case types.WorkspaceMessageTypeShellResize:
|
||||
// Resize is currently not implemented in the shell session
|
||||
// TODO: Implement shell.Resize(cols, rows) when needed
|
||||
s.logger.Debug("shell resize requested", zap.Int("cols", msg.Cols), zap.Int("rows", msg.Rows))
|
||||
case types.WorkspaceMessageTypePing:
|
||||
// Respond with pong
|
||||
pongMsg := types.NewWorkspacePong()
|
||||
if err := conn.WriteJSON(pongMsg); err != nil {
|
||||
s.logger.Debug("workspace stream pong write error", zap.Error(err))
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// Stream file updates to client
|
||||
// Forward all workspace events to WebSocket
|
||||
for {
|
||||
select {
|
||||
case update, ok := <-sub:
|
||||
case <-done:
|
||||
return
|
||||
case msg, ok := <-sub:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
data, err := json.Marshal(update)
|
||||
if err != nil {
|
||||
s.logger.Error("failed to marshal file update", zap.Error(err))
|
||||
continue
|
||||
}
|
||||
if err := conn.WriteMessage(websocket.TextMessage, data); err != nil {
|
||||
if err := conn.WriteJSON(msg); err != nil {
|
||||
s.logger.Debug("workspace stream write error", zap.Error(err))
|
||||
return
|
||||
}
|
||||
case <-closeCh:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// handleFileChangesStreamWS streams filesystem change notifications via WebSocket
|
||||
func (s *Server) handleFileChangesStreamWS(c *gin.Context) {
|
||||
conn, err := s.upgrader.Upgrade(c.Writer, c.Request, nil)
|
||||
if err != nil {
|
||||
s.logger.Error("WebSocket upgrade failed", zap.Error(err))
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
if err := conn.Close(); err != nil {
|
||||
s.logger.Debug("failed to close file changes websocket", zap.Error(err))
|
||||
}
|
||||
}()
|
||||
|
||||
// Subscribe to file change notifications
|
||||
sub := s.procMgr.GetWorkspaceTracker().SubscribeFileChanges()
|
||||
defer s.procMgr.GetWorkspaceTracker().UnsubscribeFileChanges(sub)
|
||||
|
||||
// Handle WebSocket close
|
||||
closeCh := make(chan struct{})
|
||||
go func() {
|
||||
for {
|
||||
_, _, err := conn.ReadMessage()
|
||||
if err != nil {
|
||||
close(closeCh)
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// Stream file change notifications to client
|
||||
for {
|
||||
select {
|
||||
case notification, ok := <-sub:
|
||||
case data, ok := <-shellOutputCh:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
data, err := json.Marshal(notification)
|
||||
if err != nil {
|
||||
s.logger.Error("failed to marshal file change notification", zap.Error(err))
|
||||
// Shell output channel closed
|
||||
shellOutputCh = nil
|
||||
continue
|
||||
}
|
||||
if err := conn.WriteMessage(websocket.TextMessage, data); err != nil {
|
||||
// Forward shell output as workspace stream message
|
||||
shellMsg := types.NewWorkspaceShellOutput(string(data))
|
||||
if err := conn.WriteJSON(shellMsg); err != nil {
|
||||
s.logger.Debug("workspace stream shell output write error", zap.Error(err))
|
||||
return
|
||||
}
|
||||
case <-closeCh:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -29,11 +29,9 @@ type WorkspaceTracker struct {
|
||||
currentFiles types.FileListUpdate
|
||||
mu sync.RWMutex
|
||||
|
||||
// Subscribers
|
||||
gitStatusSubscribers map[types.GitStatusSubscriber]struct{}
|
||||
filesSubscribers map[types.FilesSubscriber]struct{}
|
||||
fileChangeSubscribers map[types.FileChangeSubscriber]struct{}
|
||||
subMu sync.RWMutex
|
||||
// Unified workspace stream subscribers
|
||||
workspaceStreamSubscribers map[types.WorkspaceStreamSubscriber]struct{}
|
||||
workspaceSubMu sync.RWMutex
|
||||
|
||||
// Filesystem watcher
|
||||
watcher *fsnotify.Watcher
|
||||
@@ -55,14 +53,12 @@ func NewWorkspaceTracker(workDir string, log *logger.Logger) *WorkspaceTracker {
|
||||
}
|
||||
|
||||
return &WorkspaceTracker{
|
||||
workDir: workDir,
|
||||
logger: log.WithFields(zap.String("component", "workspace-tracker")),
|
||||
gitStatusSubscribers: make(map[types.GitStatusSubscriber]struct{}),
|
||||
filesSubscribers: make(map[types.FilesSubscriber]struct{}),
|
||||
fileChangeSubscribers: make(map[types.FileChangeSubscriber]struct{}),
|
||||
watcher: watcher,
|
||||
fsChangeTrigger: make(chan struct{}, 1), // Buffered to avoid blocking
|
||||
stopCh: make(chan struct{}),
|
||||
workDir: workDir,
|
||||
logger: log.WithFields(zap.String("component", "workspace-tracker")),
|
||||
workspaceStreamSubscribers: make(map[types.WorkspaceStreamSubscriber]struct{}),
|
||||
watcher: watcher,
|
||||
fsChangeTrigger: make(chan struct{}, 1), // Buffered to avoid blocking
|
||||
stopCh: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -147,8 +143,8 @@ func (wt *WorkspaceTracker) monitorLoop(ctx context.Context) {
|
||||
if pendingUpdate {
|
||||
wt.updateGitStatus(ctx)
|
||||
wt.updateFiles(ctx)
|
||||
// Notify file change subscribers that workspace changed
|
||||
wt.notifyFileChangeSubscribers(types.FileChangeNotification{
|
||||
// Notify workspace stream subscribers that workspace changed
|
||||
wt.notifyWorkspaceStreamFileChange(types.FileChangeNotification{
|
||||
Timestamp: time.Now(),
|
||||
Path: "",
|
||||
Operation: "refresh",
|
||||
@@ -171,8 +167,8 @@ func (wt *WorkspaceTracker) updateGitStatus(ctx context.Context) {
|
||||
wt.currentStatus = status
|
||||
wt.mu.Unlock()
|
||||
|
||||
// Notify subscribers
|
||||
wt.notifyGitStatusSubscribers(status)
|
||||
// Notify workspace stream subscribers
|
||||
wt.notifyWorkspaceStreamGitStatus(status)
|
||||
}
|
||||
|
||||
// getGitStatus retrieves the current git status
|
||||
@@ -427,8 +423,8 @@ func (wt *WorkspaceTracker) updateFiles(ctx context.Context) {
|
||||
wt.currentFiles = files
|
||||
wt.mu.Unlock()
|
||||
|
||||
// Notify subscribers
|
||||
wt.notifyFilesSubscribers(files)
|
||||
// Notify workspace stream subscribers
|
||||
wt.notifyWorkspaceStreamFileList(files)
|
||||
}
|
||||
|
||||
// getFileList retrieves the list of files in the workspace
|
||||
@@ -463,133 +459,94 @@ func (wt *WorkspaceTracker) getFileList(ctx context.Context) (types.FileListUpda
|
||||
|
||||
|
||||
|
||||
// Subscriber management methods
|
||||
// Unified workspace stream subscriber management methods
|
||||
|
||||
// SubscribeGitStatus creates a new git status subscriber
|
||||
func (wt *WorkspaceTracker) SubscribeGitStatus() types.GitStatusSubscriber {
|
||||
sub := make(types.GitStatusSubscriber, 10)
|
||||
// SubscribeWorkspaceStream creates a new unified workspace stream subscriber
|
||||
// and sends current git status and file list immediately
|
||||
func (wt *WorkspaceTracker) SubscribeWorkspaceStream() types.WorkspaceStreamSubscriber {
|
||||
sub := make(types.WorkspaceStreamSubscriber, 100)
|
||||
|
||||
wt.subMu.Lock()
|
||||
wt.gitStatusSubscribers[sub] = struct{}{}
|
||||
wt.subMu.Unlock()
|
||||
wt.workspaceSubMu.Lock()
|
||||
wt.workspaceStreamSubscribers[sub] = struct{}{}
|
||||
wt.workspaceSubMu.Unlock()
|
||||
|
||||
// Send current status immediately
|
||||
// Send current git status immediately
|
||||
wt.mu.RLock()
|
||||
current := wt.currentStatus
|
||||
currentStatus := wt.currentStatus
|
||||
currentFiles := wt.currentFiles
|
||||
wt.mu.RUnlock()
|
||||
|
||||
if current.Timestamp.IsZero() {
|
||||
current.Timestamp = time.Now()
|
||||
if currentStatus.Timestamp.IsZero() {
|
||||
currentStatus.Timestamp = time.Now()
|
||||
}
|
||||
if currentFiles.Timestamp.IsZero() {
|
||||
currentFiles.Timestamp = time.Now()
|
||||
}
|
||||
|
||||
// Send git status
|
||||
select {
|
||||
case sub <- current:
|
||||
case sub <- types.NewWorkspaceGitStatus(¤tStatus):
|
||||
default:
|
||||
}
|
||||
|
||||
// Send file list
|
||||
select {
|
||||
case sub <- types.NewWorkspaceFileList(¤tFiles):
|
||||
default:
|
||||
}
|
||||
|
||||
return sub
|
||||
}
|
||||
|
||||
// UnsubscribeGitStatus removes a git status subscriber
|
||||
func (wt *WorkspaceTracker) UnsubscribeGitStatus(sub types.GitStatusSubscriber) {
|
||||
wt.subMu.Lock()
|
||||
delete(wt.gitStatusSubscribers, sub)
|
||||
wt.subMu.Unlock()
|
||||
// UnsubscribeWorkspaceStream removes and closes a workspace stream subscriber
|
||||
func (wt *WorkspaceTracker) UnsubscribeWorkspaceStream(sub types.WorkspaceStreamSubscriber) {
|
||||
wt.workspaceSubMu.Lock()
|
||||
delete(wt.workspaceStreamSubscribers, sub)
|
||||
wt.workspaceSubMu.Unlock()
|
||||
close(sub)
|
||||
}
|
||||
|
||||
// SubscribeFiles creates a new files subscriber
|
||||
func (wt *WorkspaceTracker) SubscribeFiles() types.FilesSubscriber {
|
||||
sub := make(types.FilesSubscriber, 10)
|
||||
// Notification methods for unified workspace stream
|
||||
|
||||
wt.subMu.Lock()
|
||||
wt.filesSubscribers[sub] = struct{}{}
|
||||
wt.subMu.Unlock()
|
||||
// notifyWorkspaceStreamGitStatus sends git status to all workspace stream subscribers
|
||||
func (wt *WorkspaceTracker) notifyWorkspaceStreamGitStatus(update types.GitStatusUpdate) {
|
||||
wt.workspaceSubMu.RLock()
|
||||
defer wt.workspaceSubMu.RUnlock()
|
||||
|
||||
// Send current files immediately
|
||||
wt.mu.RLock()
|
||||
current := wt.currentFiles
|
||||
wt.mu.RUnlock()
|
||||
|
||||
if current.Timestamp.IsZero() {
|
||||
current.Timestamp = time.Now()
|
||||
}
|
||||
|
||||
select {
|
||||
case sub <- current:
|
||||
default:
|
||||
}
|
||||
|
||||
return sub
|
||||
}
|
||||
|
||||
// UnsubscribeFiles removes a files subscriber
|
||||
func (wt *WorkspaceTracker) UnsubscribeFiles(sub types.FilesSubscriber) {
|
||||
wt.subMu.Lock()
|
||||
delete(wt.filesSubscribers, sub)
|
||||
wt.subMu.Unlock()
|
||||
close(sub)
|
||||
}
|
||||
|
||||
// Notification methods
|
||||
|
||||
// notifyGitStatusSubscribers notifies all git status subscribers
|
||||
func (wt *WorkspaceTracker) notifyGitStatusSubscribers(update types.GitStatusUpdate) {
|
||||
wt.subMu.RLock()
|
||||
defer wt.subMu.RUnlock()
|
||||
|
||||
for sub := range wt.gitStatusSubscribers {
|
||||
msg := types.NewWorkspaceGitStatus(&update)
|
||||
for sub := range wt.workspaceStreamSubscribers {
|
||||
select {
|
||||
case sub <- update:
|
||||
case sub <- msg:
|
||||
default:
|
||||
// Subscriber is slow, skip
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// notifyFilesSubscribers notifies all files subscribers
|
||||
func (wt *WorkspaceTracker) notifyFilesSubscribers(update types.FileListUpdate) {
|
||||
wt.subMu.RLock()
|
||||
defer wt.subMu.RUnlock()
|
||||
// notifyWorkspaceStreamFileChange sends file change notification to all workspace stream subscribers
|
||||
func (wt *WorkspaceTracker) notifyWorkspaceStreamFileChange(notification types.FileChangeNotification) {
|
||||
wt.workspaceSubMu.RLock()
|
||||
defer wt.workspaceSubMu.RUnlock()
|
||||
|
||||
for sub := range wt.filesSubscribers {
|
||||
msg := types.NewWorkspaceFileChange(¬ification)
|
||||
for sub := range wt.workspaceStreamSubscribers {
|
||||
select {
|
||||
case sub <- update:
|
||||
case sub <- msg:
|
||||
default:
|
||||
// Subscriber is slow, skip
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// notifyWorkspaceStreamFileList sends file list update to all workspace stream subscribers
|
||||
func (wt *WorkspaceTracker) notifyWorkspaceStreamFileList(update types.FileListUpdate) {
|
||||
wt.workspaceSubMu.RLock()
|
||||
defer wt.workspaceSubMu.RUnlock()
|
||||
|
||||
|
||||
// SubscribeFileChanges creates a new file change subscriber
|
||||
func (wt *WorkspaceTracker) SubscribeFileChanges() types.FileChangeSubscriber {
|
||||
sub := make(types.FileChangeSubscriber, 100)
|
||||
|
||||
wt.subMu.Lock()
|
||||
wt.fileChangeSubscribers[sub] = struct{}{}
|
||||
wt.subMu.Unlock()
|
||||
|
||||
return sub
|
||||
}
|
||||
|
||||
// UnsubscribeFileChanges removes a file change subscriber
|
||||
func (wt *WorkspaceTracker) UnsubscribeFileChanges(sub types.FileChangeSubscriber) {
|
||||
wt.subMu.Lock()
|
||||
delete(wt.fileChangeSubscribers, sub)
|
||||
wt.subMu.Unlock()
|
||||
close(sub)
|
||||
}
|
||||
|
||||
// notifyFileChangeSubscribers notifies all file change subscribers
|
||||
func (wt *WorkspaceTracker) notifyFileChangeSubscribers(notification types.FileChangeNotification) {
|
||||
wt.subMu.RLock()
|
||||
defer wt.subMu.RUnlock()
|
||||
|
||||
for sub := range wt.fileChangeSubscribers {
|
||||
msg := types.NewWorkspaceFileList(&update)
|
||||
for sub := range wt.workspaceStreamSubscribers {
|
||||
select {
|
||||
case sub <- notification:
|
||||
case sub <- msg:
|
||||
default:
|
||||
// Subscriber is slow, skip
|
||||
}
|
||||
|
||||
@@ -27,24 +27,28 @@
|
||||
//
|
||||
// Message type: PermissionNotification
|
||||
//
|
||||
// # Git Status Stream (/api/v1/workspace/git-status/stream)
|
||||
// # Unified Workspace Stream (/api/v1/workspace/stream)
|
||||
//
|
||||
// Streams git status updates when the workspace state changes.
|
||||
// Bidirectional WebSocket that consolidates all workspace-related streams:
|
||||
// - Shell I/O (input, output, exit, resize)
|
||||
// - Git status updates
|
||||
// - File change notifications
|
||||
// - File list updates
|
||||
// - Ping/pong keepalive
|
||||
//
|
||||
// Message type: GitStatusUpdate
|
||||
// Message type: WorkspaceStreamMessage (defined in types/types.go)
|
||||
//
|
||||
// # File Changes Stream (/api/v1/workspace/file-changes/stream)
|
||||
//
|
||||
// Streams file system change notifications when files are created,
|
||||
// modified, deleted, or renamed in the workspace.
|
||||
//
|
||||
// Message type: FileChangeNotification
|
||||
//
|
||||
// # Shell Stream (/api/v1/shell/stream)
|
||||
//
|
||||
// Bidirectional WebSocket for interactive shell I/O.
|
||||
//
|
||||
// Message type: ShellMessage (input/output/ping/pong/exit)
|
||||
// Message types (use WorkspaceMsg* constants):
|
||||
// - shell_output: Shell output data
|
||||
// - shell_input: Shell input data (client -> server)
|
||||
// - shell_exit: Shell process exited
|
||||
// - shell_resize: Terminal resize (client -> server)
|
||||
// - git_status: Git status update
|
||||
// - file_change: File change notification
|
||||
// - file_list: File list update
|
||||
// - ping/pong: Keepalive messages
|
||||
// - connected: Connection established
|
||||
// - error: Error occurred
|
||||
//
|
||||
// All streams use JSON-encoded messages over WebSocket connections.
|
||||
package streams
|
||||
|
||||
@@ -6,6 +6,7 @@ package types
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/kandev/kandev/internal/agentctl/types/streams"
|
||||
)
|
||||
@@ -111,14 +112,135 @@ type PermissionResponse struct {
|
||||
// PermissionHandler is called when the agent requests permission for an action.
|
||||
type PermissionHandler func(ctx context.Context, req *PermissionRequest) (*PermissionResponse, error)
|
||||
|
||||
// Subscriber types for internal use.
|
||||
type (
|
||||
// GitStatusSubscriber is a channel that receives git status updates.
|
||||
GitStatusSubscriber chan GitStatusUpdate
|
||||
// WorkspaceMessageType represents the type of workspace stream message
|
||||
type WorkspaceMessageType string
|
||||
|
||||
// FilesSubscriber is a channel that receives file listing updates.
|
||||
FilesSubscriber chan FileListUpdate
|
||||
|
||||
// FileChangeSubscriber is a channel that receives file change notifications.
|
||||
FileChangeSubscriber chan FileChangeNotification
|
||||
const (
|
||||
// Workspace stream message types
|
||||
WorkspaceMessageTypeShellOutput WorkspaceMessageType = "shell_output"
|
||||
WorkspaceMessageTypeShellInput WorkspaceMessageType = "shell_input"
|
||||
WorkspaceMessageTypeShellExit WorkspaceMessageType = "shell_exit"
|
||||
WorkspaceMessageTypePing WorkspaceMessageType = "ping"
|
||||
WorkspaceMessageTypePong WorkspaceMessageType = "pong"
|
||||
WorkspaceMessageTypeGitStatus WorkspaceMessageType = "git_status"
|
||||
WorkspaceMessageTypeFileChange WorkspaceMessageType = "file_change"
|
||||
WorkspaceMessageTypeFileList WorkspaceMessageType = "file_list"
|
||||
WorkspaceMessageTypeError WorkspaceMessageType = "error"
|
||||
WorkspaceMessageTypeConnected WorkspaceMessageType = "connected"
|
||||
WorkspaceMessageTypeShellResize WorkspaceMessageType = "shell_resize"
|
||||
)
|
||||
|
||||
// WorkspaceStreamMessage is the unified message format for the workspace stream.
|
||||
// It carries all workspace events (shell I/O, git status, file changes) with
|
||||
// message type differentiation.
|
||||
type WorkspaceStreamMessage struct {
|
||||
Type WorkspaceMessageType `json:"type"`
|
||||
Timestamp int64 `json:"timestamp"` // Unix milliseconds
|
||||
|
||||
// Shell fields (for shell_output, shell_input, shell_exit)
|
||||
Data string `json:"data,omitempty"` // Shell output or input data
|
||||
Code int `json:"code,omitempty"` // Exit code for shell_exit
|
||||
|
||||
// Shell resize fields (for shell_resize)
|
||||
Cols int `json:"cols,omitempty"`
|
||||
Rows int `json:"rows,omitempty"`
|
||||
|
||||
// Git status fields (for git_status)
|
||||
GitStatus *GitStatusUpdate `json:"git_status,omitempty"`
|
||||
|
||||
// File change fields (for file_change)
|
||||
FileChange *FileChangeNotification `json:"file_change,omitempty"`
|
||||
|
||||
// File list fields (for file_list)
|
||||
FileList *FileListUpdate `json:"file_list,omitempty"`
|
||||
|
||||
// Error fields (for error)
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// NewWorkspaceShellOutput creates a shell output message
|
||||
func NewWorkspaceShellOutput(data string) WorkspaceStreamMessage {
|
||||
return WorkspaceStreamMessage{
|
||||
Type: WorkspaceMessageTypeShellOutput,
|
||||
Timestamp: timeNowUnixMilli(),
|
||||
Data: data,
|
||||
}
|
||||
}
|
||||
|
||||
// NewWorkspaceShellInput creates a shell input message
|
||||
func NewWorkspaceShellInput(data string) WorkspaceStreamMessage {
|
||||
return WorkspaceStreamMessage{
|
||||
Type: WorkspaceMessageTypeShellInput,
|
||||
Timestamp: timeNowUnixMilli(),
|
||||
Data: data,
|
||||
}
|
||||
}
|
||||
|
||||
// NewWorkspaceGitStatus creates a git status message
|
||||
func NewWorkspaceGitStatus(status *GitStatusUpdate) WorkspaceStreamMessage {
|
||||
return WorkspaceStreamMessage{
|
||||
Type: WorkspaceMessageTypeGitStatus,
|
||||
Timestamp: timeNowUnixMilli(),
|
||||
GitStatus: status,
|
||||
}
|
||||
}
|
||||
|
||||
// NewWorkspaceFileChange creates a file change message
|
||||
func NewWorkspaceFileChange(notification *FileChangeNotification) WorkspaceStreamMessage {
|
||||
return WorkspaceStreamMessage{
|
||||
Type: WorkspaceMessageTypeFileChange,
|
||||
Timestamp: timeNowUnixMilli(),
|
||||
FileChange: notification,
|
||||
}
|
||||
}
|
||||
|
||||
// NewWorkspaceFileList creates a file list message
|
||||
func NewWorkspaceFileList(update *FileListUpdate) WorkspaceStreamMessage {
|
||||
return WorkspaceStreamMessage{
|
||||
Type: WorkspaceMessageTypeFileList,
|
||||
Timestamp: timeNowUnixMilli(),
|
||||
FileList: update,
|
||||
}
|
||||
}
|
||||
|
||||
// NewWorkspacePong creates a pong message
|
||||
func NewWorkspacePong() WorkspaceStreamMessage {
|
||||
return WorkspaceStreamMessage{
|
||||
Type: WorkspaceMessageTypePong,
|
||||
Timestamp: timeNowUnixMilli(),
|
||||
}
|
||||
}
|
||||
|
||||
// NewWorkspaceConnected creates a connected message
|
||||
func NewWorkspaceConnected() WorkspaceStreamMessage {
|
||||
return WorkspaceStreamMessage{
|
||||
Type: WorkspaceMessageTypeConnected,
|
||||
Timestamp: timeNowUnixMilli(),
|
||||
}
|
||||
}
|
||||
|
||||
// NewWorkspacePing creates a ping message
|
||||
func NewWorkspacePing() WorkspaceStreamMessage {
|
||||
return WorkspaceStreamMessage{
|
||||
Type: WorkspaceMessageTypePing,
|
||||
Timestamp: timeNowUnixMilli(),
|
||||
}
|
||||
}
|
||||
|
||||
// NewWorkspaceShellResize creates a shell resize message
|
||||
func NewWorkspaceShellResize(cols, rows int) WorkspaceStreamMessage {
|
||||
return WorkspaceStreamMessage{
|
||||
Type: WorkspaceMessageTypeShellResize,
|
||||
Timestamp: timeNowUnixMilli(),
|
||||
Cols: cols,
|
||||
Rows: rows,
|
||||
}
|
||||
}
|
||||
|
||||
// WorkspaceStreamSubscriber is a channel that receives unified workspace messages
|
||||
type WorkspaceStreamSubscriber chan WorkspaceStreamMessage
|
||||
|
||||
// timeNowUnixMilli returns current time in unix milliseconds
|
||||
func timeNowUnixMilli() int64 {
|
||||
return time.Now().UnixMilli()
|
||||
}
|
||||
|
||||
@@ -104,6 +104,32 @@ const (
|
||||
FileChangeNotified = "file.change.notified" // File changed in workspace
|
||||
)
|
||||
|
||||
// Event types for shell I/O
|
||||
const (
|
||||
ShellOutput = "shell.output" // Shell output data
|
||||
ShellExit = "shell.exit" // Shell process exited
|
||||
)
|
||||
|
||||
// BuildShellOutputSubject creates a shell output subject for a specific session
|
||||
func BuildShellOutputSubject(sessionID string) string {
|
||||
return ShellOutput + "." + sessionID
|
||||
}
|
||||
|
||||
// BuildShellOutputWildcardSubject creates a wildcard subscription for all shell output events
|
||||
func BuildShellOutputWildcardSubject() string {
|
||||
return ShellOutput + ".*"
|
||||
}
|
||||
|
||||
// BuildShellExitSubject creates a shell exit subject for a specific session
|
||||
func BuildShellExitSubject(sessionID string) string {
|
||||
return ShellExit + "." + sessionID
|
||||
}
|
||||
|
||||
// BuildShellExitWildcardSubject creates a wildcard subscription for all shell exit events
|
||||
func BuildShellExitWildcardSubject() string {
|
||||
return ShellExit + ".*"
|
||||
}
|
||||
|
||||
// BuildAgentStreamSubject creates an agent stream subject for a specific session
|
||||
func BuildAgentStreamSubject(sessionID string) string {
|
||||
return AgentStream + "." + sessionID
|
||||
|
||||
@@ -159,10 +159,10 @@ const (
|
||||
ActionWorkspaceFileChanges = "session.workspace.file.changes" // Notification
|
||||
|
||||
// Shell actions
|
||||
ActionShellStatus = "session.shell.status" // Get shell status
|
||||
ActionShellSubscribe = "shell.subscribe" // Subscribe to shell output
|
||||
ActionShellInput = "shell.input" // Send input to shell
|
||||
ActionShellOutput = "session.shell.output" // Shell output notification
|
||||
ActionShellStatus = "session.shell.status" // Get shell status
|
||||
ActionShellSubscribe = "shell.subscribe" // Subscribe to shell output
|
||||
ActionShellInput = "shell.input" // Send input to shell
|
||||
ActionSessionShellOutput = "session.shell.output" // Shell output notification (also used for exit with type: "exit")
|
||||
|
||||
// User actions
|
||||
ActionUserGet = "user.get"
|
||||
|
||||
@@ -5,7 +5,6 @@ import { IconChevronRight, IconChevronDown, IconFile, IconFolder, IconFolderOpen
|
||||
import { ScrollArea } from '@kandev/ui/scroll-area';
|
||||
import { getWebSocketClient } from '@/lib/ws/connection';
|
||||
import { requestFileTree, requestFileContent } from '@/lib/ws/workspace-files';
|
||||
import { useSessionAgentctl } from '@/hooks/use-session-agentctl';
|
||||
import type { FileTreeNode, FileContentResponse, OpenFileTab } from '@/lib/types/backend';
|
||||
|
||||
type FileBrowserProps = {
|
||||
@@ -14,21 +13,18 @@ type FileBrowserProps = {
|
||||
};
|
||||
|
||||
export function FileBrowser({ sessionId, onOpenFile }: FileBrowserProps) {
|
||||
const { isReady } = useSessionAgentctl(sessionId);
|
||||
const [tree, setTree] = useState<FileTreeNode | null>(null);
|
||||
const [expandedPaths, setExpandedPaths] = useState<Set<string>>(new Set());
|
||||
const [loadingPaths, setLoadingPaths] = useState<Set<string>>(new Set());
|
||||
const [isLoadingTree, setIsLoadingTree] = useState(true);
|
||||
|
||||
// Load initial tree
|
||||
// Load initial tree - always try to load, don't wait for agentctl ready status
|
||||
useEffect(() => {
|
||||
setTree(null);
|
||||
setExpandedPaths(new Set());
|
||||
setLoadingPaths(new Set());
|
||||
setIsLoadingTree(true);
|
||||
if (!isReady) {
|
||||
return;
|
||||
}
|
||||
|
||||
const loadTree = async () => {
|
||||
try {
|
||||
const client = getWebSocketClient();
|
||||
@@ -44,11 +40,10 @@ export function FileBrowser({ sessionId, onOpenFile }: FileBrowserProps) {
|
||||
};
|
||||
|
||||
loadTree();
|
||||
}, [isReady, sessionId]);
|
||||
}, [sessionId]);
|
||||
|
||||
// Subscribe to file changes and refresh tree
|
||||
useEffect(() => {
|
||||
if (!isReady) return;
|
||||
const client = getWebSocketClient();
|
||||
if (!client) return;
|
||||
|
||||
@@ -63,7 +58,7 @@ export function FileBrowser({ sessionId, onOpenFile }: FileBrowserProps) {
|
||||
});
|
||||
|
||||
return unsubscribe;
|
||||
}, [isReady, sessionId]);
|
||||
}, [sessionId]);
|
||||
|
||||
const toggleExpand = useCallback(async (node: FileTreeNode) => {
|
||||
if (!node.is_dir) return;
|
||||
@@ -116,7 +111,6 @@ export function FileBrowser({ sessionId, onOpenFile }: FileBrowserProps) {
|
||||
const openFile = useCallback(async (node: FileTreeNode) => {
|
||||
if (node.is_dir) return;
|
||||
|
||||
if (!isReady) return;
|
||||
try {
|
||||
const client = getWebSocketClient();
|
||||
if (!client) return;
|
||||
@@ -131,7 +125,7 @@ export function FileBrowser({ sessionId, onOpenFile }: FileBrowserProps) {
|
||||
} catch (error) {
|
||||
console.error('Failed to load file content:', error);
|
||||
}
|
||||
}, [isReady, sessionId, onOpenFile]);
|
||||
}, [sessionId, onOpenFile]);
|
||||
|
||||
const renderTreeNode = (node: FileTreeNode, depth: number = 0): React.ReactNode => {
|
||||
const isExpanded = expandedPaths.has(node.path);
|
||||
|
||||
Reference in New Issue
Block a user