// Package copilot provides a Go SDK for interacting with the GitHub Copilot CLI.

package copilot

import (

"context"

"encoding/json"

"errors"

"fmt"

"log"

"strings"

"sync"

"sync/atomic"

"time"

"github.com/github/copilot-sdk/go/internal/jsonrpc2"

"github.com/github/copilot-sdk/go/rpc"

)

// toolSearchToolName is the fixed name of the runtime's built-in tool-search

// tool. A client can replace its behavior by registering a [Tool] with this

// exact name and OverridesBuiltInTool set to true.

const toolSearchToolName = "tool_search_tool"

type sessionHandler struct {

id uint64

fn SessionEventHandler

}

// Session represents a single conversation session with the Copilot CLI.

//

// A session maintains conversation state, handles events, and manages tool execution.

// Sessions are created via [Client.CreateSession] or resumed via [Client.ResumeSession].

//

// The session provides methods to send messages, subscribe to events, retrieve

// conversation history, and manage the session lifecycle. All methods are safe

// for concurrent use.

//

// Example usage:

//

// session, err := client.CreateSession(copilot.SessionConfig{

// Model: "gpt-4",

// })

// if err != nil {

// log.Fatal(err)

// }

// defer session.Disconnect()

//

// // Subscribe to events

// unsubscribe := session.On(func(event copilot.SessionEvent) {

// if d, ok := event.Data.(*copilot.AssistantMessageData); ok {

// fmt.Println("Assistant:", d.Content)

// }

// })

// defer unsubscribe()

//

// // Send a message

// messageID, err := session.Send(copilot.MessageOptions{

// Prompt: "Hello, world!",

// })

type Session struct {

// SessionID is the unique identifier for this session.

SessionID string

workspacePath string

transcriptRecovery *TranscriptRecoveryReport

client *jsonrpc2.Client

clientSessionAPIs *rpc.ClientSessionAPIHandlers

handlers []*sessionHandler

nextHandlerID uint64

handlerMutex sync.RWMutex

gracefulRetirement bool

toolHandlers map[string]ToolHandler

toolHandlersM sync.RWMutex

setToolsOnce sync.Once

setToolsGate chan struct{}

pendingExternalTools map[string]*pendingExternalTool

pendingExternalToolsM sync.Mutex

externalToolsClosed bool

permissionHandler permissionHandlerFunc

permissionMux sync.RWMutex

managedSettings bool

mcpAuthHandler MCPAuthHandler

mcpAuthRefreshHandler MCPAuthRuntimeRefreshHandler

mcpAuthMu sync.RWMutex

userInputHandler UserInputHandler

userInputMux sync.RWMutex

exitPlanModeHandler ExitPlanModeRequestHandler

exitPlanModeMu sync.RWMutex

autoModeSwitchHandler AutoModeSwitchRequestHandler

autoModeSwitchMu sync.RWMutex

hooks *SessionHooks

hooksMux sync.RWMutex

transformCallbacks map[string]SectionTransformFn

transformMu sync.Mutex

commandHandlers map[string]CommandHandler

commandHandlersMu sync.RWMutex

elicitationHandler ElicitationHandler

elicitationMu sync.RWMutex

canvasHandler CanvasHandler

canvasMu sync.RWMutex

skillProvider SkillProvider

skillProviderMu sync.RWMutex

bearerTokenProviders map[string]BearerTokenProvider

bearerTokenMu sync.RWMutex

releaseGitHubTokenProvider func()

gitHubTokenProviderMu sync.Mutex

gitHubTokenProviderReleased bool

openCanvases []rpc.OpenCanvasInstance

openCanvasesMu sync.RWMutex

capabilities SessionCapabilities

capabilitiesMu sync.RWMutex

// eventCh serializes user event handler dispatch. dispatchEvent enqueues;

// a single goroutine (processEvents) dequeues and invokes handlers in FIFO order.

// eventDone stops both sides without closing eventCh while a sender may be active.

eventCh chan SessionEvent

eventDone chan struct{}

closeOnce sync.Once

completionRPCUnsupported atomic.Bool

// RPC provides typed session-scoped RPC methods.

RPC *rpc.SessionRPC

}

type pendingExternalTool struct {

ctx context.Context

cancel context.CancelFunc

}

// WorkspacePath returns the path to the session workspace directory when infinite

// sessions are enabled. Contains checkpoints/, plan.md, and files/ subdirectories.

// Returns empty string if infinite sessions are disabled.

func (s *Session) WorkspacePath() string {

return s.workspacePath

}

// TranscriptRecovery returns the report from session.resume, or nil if no

// transcript repair was reported. The returned value is independent of session state.

func (s *Session) TranscriptRecovery() *TranscriptRecoveryReport {

if s.transcriptRecovery == nil {

return nil

}

report := *s.transcriptRecovery

report.InvalidLineNumbers = append([]int(nil), report.InvalidLineNumbers...)

return &report

}

// OpenCanvases returns the open-canvas snapshot last reported by the runtime.

// The snapshot is populated from session.resume and live session.canvas.opened

// and session.canvas.closed events. The returned slice is a copy and is safe to

// mutate by the caller.

func (s *Session) OpenCanvases() []rpc.OpenCanvasInstance {

s.openCanvasesMu.RLock()

defer s.openCanvasesMu.RUnlock()

if len(s.openCanvases) == 0 {

return nil

}

out := make([]rpc.OpenCanvasInstance, len(s.openCanvases))

copy(out, s.openCanvases)

return out

}

func (s *Session) setOpenCanvases(canvases []rpc.OpenCanvasInstance) {

s.openCanvasesMu.Lock()

defer s.openCanvasesMu.Unlock()

s.openCanvases = canvases

}

func (s *Session) upsertOpenCanvas(canvas rpc.OpenCanvasInstance) {

s.openCanvasesMu.Lock()

defer s.openCanvasesMu.Unlock()

for i := range s.openCanvases {

if s.openCanvases[i].InstanceID == canvas.InstanceID {

s.openCanvases[i] = canvas

return

}

}

s.openCanvases = append(s.openCanvases, canvas)

}

func (s *Session) removeOpenCanvas(instanceID string) {

s.openCanvasesMu.Lock()

defer s.openCanvasesMu.Unlock()

filtered := make([]rpc.OpenCanvasInstance, 0, len(s.openCanvases))

for _, canvas := range s.openCanvases {

if canvas.InstanceID != instanceID {

filtered = append(filtered, canvas)

}

}

s.openCanvases = filtered

}

func (s *Session) updateOpenCanvasesFromEvent(event SessionEvent) {

switch data := event.Data.(type) {

case *SessionCanvasOpenedData:

if data.InstanceID == "" || data.CanvasID == "" || data.ExtensionID == "" {

fmt.Printf("failed to deserialize session.canvas.opened payload\n")

return

}

s.upsertOpenCanvas(rpc.OpenCanvasInstance{

CanvasID: data.CanvasID,

ExtensionID: data.ExtensionID,

ExtensionName: data.ExtensionName,

Input: data.Input,

InstanceID: data.InstanceID,

Status: data.Status,

Title: data.Title,

Icon: data.Icon,

URL: data.URL,

})

case *SessionCanvasClosedData:

if data.InstanceID == "" {

fmt.Printf("failed to deserialize session.canvas.closed payload\n")

return

}

s.removeOpenCanvas(data.InstanceID)

}

}

func (s *Session) registerCanvasHandler(handler CanvasHandler) {

s.canvasMu.Lock()

defer s.canvasMu.Unlock()

s.canvasHandler = handler

}

func (s *Session) getCanvasHandler() CanvasHandler {

s.canvasMu.RLock()

defer s.canvasMu.RUnlock()

return s.canvasHandler

}

func (s *Session) registerSkillProvider(provider SkillProvider) {

s.skillProviderMu.Lock()

defer s.skillProviderMu.Unlock()

s.skillProvider = provider

}

func (s *Session) getSkillProvider() SkillProvider {

s.skillProviderMu.RLock()

defer s.skillProviderMu.RUnlock()

return s.skillProvider

}

func (s *Session) clearSkillProvider() {

s.skillProviderMu.Lock()

s.skillProvider = nil

s.skillProviderMu.Unlock()

}

// registerBearerTokenProviders installs per-provider [BearerTokenProvider] callbacks

// for BYOK providers configured with managed-identity / on-demand bearer-token

// auth, keyed by provider name.

//

// The runtime never receives the callback itself; the SDK strips it from the

// provider config and instead sends `hasBearerTokenProvider: true`. When the

// runtime needs a token it issues a session-scoped `providerToken.getToken`

// request, which the session's provider-token adapter routes to the matching

// per-provider callback.

func (s *Session) registerBearerTokenProviders(providers map[string]BearerTokenProvider) {

s.bearerTokenMu.Lock()

defer s.bearerTokenMu.Unlock()

s.bearerTokenProviders = make(map[string]BearerTokenProvider, len(providers))

for name, callback := range providers {

if callback == nil {

continue

}

s.bearerTokenProviders[name] = callback

}

}

func (s *Session) getBearerTokenProvider(providerName string) BearerTokenProvider {

s.bearerTokenMu.RLock()

defer s.bearerTokenMu.RUnlock()

return s.bearerTokenProviders[providerName]

}

type providerTokenClientSessionAdapter struct {

session *Session

}

func newProviderTokenClientSessionAdapter(session *Session) rpc.ProviderTokenHandler {

return &providerTokenClientSessionAdapter{session: session}

}

func (a *providerTokenClientSessionAdapter) GetToken(request *rpc.ProviderTokenAcquireRequest) (*rpc.ProviderTokenAcquireResult, error) {

if request == nil {

return nil, providerTokenJSONRPCError("missing provider token request")

}

if a.session == nil || a.session.SessionID != request.SessionID {

return nil, providerTokenJSONRPCError(fmt.Sprintf("unknown session %s", request.SessionID))

}

callback := a.session.getBearerTokenProvider(request.ProviderName)

if callback == nil {

return nil, providerTokenJSONRPCError(fmt.Sprintf("No bearer-token provider registered for provider %q", request.ProviderName))

}

token, err := callback(ProviderTokenArgs{ProviderName: request.ProviderName, SessionID: request.SessionID})

if err != nil {

return nil, providerTokenJSONRPCError(err.Error())

}

return &rpc.ProviderTokenAcquireResult{Token: token}, nil

}

func providerTokenJSONRPCError(message string) *jsonrpc2.Error {

return &jsonrpc2.Error{

Code: -32603,

Message: message,

}

}

type canvasClientSessionAdapter struct {

session *Session

}

func newCanvasClientSessionAdapter(session *Session) rpc.CanvasHandler {

return &canvasClientSessionAdapter{session: session}

}

func (a *canvasClientSessionAdapter) Close(request *rpc.CanvasProviderCloseRequest) (*rpc.CanvasCloseResult, error) {

if request == nil {

return nil, canvasJSONRPCError(NewCanvasError("canvas_handler_unset", "missing canvas close request"))

}

handler, err := a.resolveHandler(canvasProviderSessionID(request))

if err != nil {

return nil, err

}

if err := handler.OnClose(context.Background(), *request); err != nil {

return nil, canvasResultError(err)

}

return nil, nil

}

func (a *canvasClientSessionAdapter) Invoke(request *rpc.CanvasProviderInvokeActionRequest) (any, error) {

if request == nil {

return nil, canvasJSONRPCError(NewCanvasError("canvas_handler_unset", "missing canvas action request"))

}

handler, err := a.resolveHandler(canvasProviderSessionID(request))

if err != nil {

return nil, err

}

result, actionErr := handler.OnAction(context.Background(), *request)

if actionErr != nil {

return nil, canvasResultError(actionErr)

}

return result, nil

}

func (a *canvasClientSessionAdapter) Open(request *rpc.CanvasProviderOpenRequest) (*rpc.CanvasProviderOpenResult, error) {

if request == nil {

return nil, canvasJSONRPCError(NewCanvasError("canvas_handler_unset", "missing canvas open request"))

}

handler, err := a.resolveHandler(canvasProviderSessionID(request))

if err != nil {

return nil, err

}

result, openErr := handler.OnOpen(context.Background(), *request)

if openErr != nil {

return nil, canvasResultError(openErr)

}

return &result, nil

}

func (a *canvasClientSessionAdapter) resolveHandler(sessionID string) (CanvasHandler, error) {

if sessionID == "" {

return nil, canvasJSONRPCError(NewCanvasError("canvas_handler_unset", "missing session ID"))

}

if a.session == nil || a.session.SessionID != sessionID {

return nil, canvasJSONRPCError(NewCanvasError("canvas_handler_unset", fmt.Sprintf("unknown session %s", sessionID)))

}

handler := a.session.getCanvasHandler()

if handler == nil {

return nil, canvasJSONRPCError(NewCanvasError(

"canvas_handler_unset",

"No CanvasHandler installed on this session; install one via SessionConfig.CanvasHandler before creating the session.",

))

}

return handler, nil

}

func canvasProviderSessionID(request any) string {

switch req := request.(type) {

case *rpc.CanvasProviderCloseRequest:

if req != nil {

return req.SessionID

}

case *rpc.CanvasProviderInvokeActionRequest:

if req != nil {

return req.SessionID

}

case *rpc.CanvasProviderOpenRequest:

if req != nil {

return req.SessionID

}

}

return ""

}

func canvasJSONRPCError(cerr *CanvasError) *jsonrpc2.Error {

data, _ := json.Marshal(map[string]string{

"code": cerr.Code,

"message": cerr.Message,

})

return &jsonrpc2.Error{

Code: -32603,

Message: cerr.Message,

Data: data,

}

}

func canvasResultError(err error) error {

if err == nil {

return nil

}

if rpcErr, ok := err.(*jsonrpc2.Error); ok {

return rpcErr

}

if cerr, ok := err.(*CanvasError); ok {

return canvasJSONRPCError(cerr)

}

return canvasJSONRPCError(NewCanvasError("canvas_handler_error", err.Error()))

}

// newSession creates a new session wrapper with the given session ID and client.

func newSession(

sessionID string,

client *jsonrpc2.Client,

workspacePath string,

managedSettings bool,

) *Session {

s := &Session{

SessionID: sessionID,

workspacePath: workspacePath,

managedSettings: managedSettings,

client: client,

clientSessionAPIs: &rpc.ClientSessionAPIHandlers{},

handlers: make([]*sessionHandler, 0),

toolHandlers: make(map[string]ToolHandler),

commandHandlers: make(map[string]CommandHandler),

eventCh: make(chan SessionEvent, 128),

eventDone: make(chan struct{}),

RPC: rpc.NewSessionRPC(client, sessionID),

}

s.clientSessionAPIs.Canvas = newCanvasClientSessionAdapter(s)

s.clientSessionAPIs.ProviderToken = newProviderTokenClientSessionAdapter(s)

go s.processEvents()

return s

}

// Send sends a message to this session and waits for the response.

//

// The message is processed asynchronously. Subscribe to events via [Session.On]

// to receive streaming responses and other session events.

//

// Parameters:

// - options: The message options including the prompt and optional attachments.

//

// Returns the message ID of the response, which can be used to correlate events,

// or an error if the session has been disconnected or the connection fails.

//

// Example:

//

// messageID, err := session.Send(context.Background(), copilot.MessageOptions{

// Prompt: "Explain this code",

// Attachments: []copilot.Attachment{

// &copilot.AttachmentFile{DisplayName: "main.go", Path: "./main.go"},

// },

// })

// if err != nil {

// log.Printf("Failed to send message: %v", err)

// }

func (s *Session) Send(ctx context.Context, options MessageOptions) (string, error) {

result, err := s.client.Request(ctx, "session.send", s.createSendRequest(ctx, options))

if err != nil {

return "", fmt.Errorf("failed to send message: %w", err)

}

var response sessionSendResponse

if err := json.Unmarshal(result, &response); err != nil {

return "", fmt.Errorf("failed to unmarshal send response: %w", err)

}

return response.MessageID, nil

}

func (s *Session) createSendRequest(ctx context.Context, options MessageOptions) sessionSendRequest {

traceparent, tracestate := getTraceContext(ctx)

req := sessionSendRequest{

SessionID: s.SessionID,

Prompt: options.Prompt,

Source: options.Source,

DisplayPrompt: options.DisplayPrompt,

Attachments: options.Attachments,

Mode: options.Mode,

AgentMode: options.AgentMode,

Traceparent: traceparent,

Tracestate: tracestate,

RequestHeaders: options.RequestHeaders,

}

if options.ResponseSchema != nil {

strict := true

req.ResponseFormat = &rpc.ResponseFormat{

Type: "json_schema",

JSONSchema: rpc.JSONSchemaResponseFormat{

Name: "response", Schema: options.ResponseSchema, Strict: &strict,

},

}

}

return req

}

// SendPrompt is a convenience wrapper for [Session.Send] that takes a plain

// prompt string instead of a [MessageOptions] struct. Equivalent to:

//

// session.Send(ctx, copilot.MessageOptions{Prompt: prompt})

func (s *Session) SendPrompt(ctx context.Context, prompt string) (string, error) {

return s.Send(ctx, MessageOptions{Prompt: prompt})

}

// SendAndWait sends a message and waits for its response to complete.

//

// Supported local runtimes correlate completion and the response with this

// message. Unsupported runtimes retain the legacy non-Autopilot session.idle or

// request-correlated session.autopilot_stopped wait. Autopilot completion does

// not wait for or stop attached shells.

//

// Events are still delivered to handlers registered via [Session.On] while waiting.

// Synchronous handlers registered before this call process the terminal event,

// when present, before it returns successfully, unless the session retires.

// This does not wait for asynchronous work started by a handler.

// Sub-agent events with a non-empty AgentID do not complete the wait or supply its reply.

//

// Parameters:

// - options: The message options including the prompt and optional attachments.

// - timeout: How long to wait for completion. Defaults to 60 seconds if zero.

// Controls how long to wait; does not abort in-flight agent work.

//

// Returns the final assistant message event, or nil if none was received.

// Returns an error if the request is aborted, the timeout is reached, the

// session is closed, or the connection fails.

//

// Example:

//

// response, err := session.SendAndWait(context.Background(), copilot.MessageOptions{

// Prompt: "What is 2+2?",

// }) // Use default 60s timeout

// if err != nil {

// log.Printf("Failed: %v", err)

// }

// if response != nil {

// if d, ok := response.Data.(*AssistantMessageData); ok {

// fmt.Println(d.Content)

// }

// }

func (s *Session) SendAndWait(ctx context.Context, options MessageOptions) (*SessionEvent, error) {

if _, ok := ctx.Deadline(); !ok {

var cancel context.CancelFunc

ctx, cancel = context.WithTimeout(ctx, 60*time.Second)

defer cancel()

}

if event, supported, err := s.sendWithCompletion(ctx, options); supported {

return event, err

}

if options.ResponseSchema != nil {

return s.sendAndWaitStructured(ctx, options)

}

return s.sendAndWaitLegacy(ctx, options)

}

func (s *Session) sendWithCompletion(ctx context.Context, options MessageOptions) (*SessionEvent, bool, error) {

if s.completionRPCUnsupported.Load() {

return nil, false, nil

}

requestCtx, cancel := context.WithCancelCause(ctx)

defer cancel(nil)

go func() {

select {

case <-s.eventDone:

cancel(errors.New("session closed before response completed"))

case <-requestCtx.Done():

}

}()

var mu sync.Mutex

observed := make(map[string]bool)

dispatched := make(map[string]bool)

boundary := ""

drained := make(chan struct{})

terminal := func(event SessionEvent) bool {

if event.AgentID != nil && *event.AgentID != "" {

return false

}

switch event.Data.(type) {

case *SessionIdleData, *SessionAutopilotStoppedData, *SessionErrorData:

return true

}

return false

}

unobserve := s.on(func(event SessionEvent) {

if terminal(event) {

mu.Lock()

observed[event.ID] = true

mu.Unlock()

}

}, true)

defer unobserve()

unsubscribe := s.On(func(event SessionEvent) {

if terminal(event) {

mu.Lock()

dispatched[event.ID] = true

if boundary == event.ID {

close(drained)

boundary = ""

}

mu.Unlock()

}

})

defer unsubscribe()

drain := func(id *string) error {

if id == nil {

return nil

}

mu.Lock()

done := dispatched[*id]

boundary = *id

mu.Unlock()

if done {

return nil

}

select {

case <-drained:

return nil

case <-s.eventDone:

s.handlerMutex.RLock()

graceful := s.gracefulRetirement

s.handlerMutex.RUnlock()

mu.Lock()

done := dispatched[*id] || (graceful && observed[*id])

mu.Unlock()

if done {

return nil

}

return errors.New("session closed before response completed")

case <-ctx.Done():

return context.Cause(ctx)

}

}

result, err := s.client.Request(requestCtx, "session.sendAndWait", s.createSendRequest(ctx, options))

if err != nil {

var rpcError *jsonrpc2.Error

if errors.As(err, &rpcError) {

if rpcError.Code == -32601 {

s.completionRPCUnsupported.Store(true)

select {

case <-s.eventDone:

return nil, true, errors.New("session closed before response completed")

default:

}

return nil, false, nil

}

if len(rpcError.Data) > 0 {

var data struct {

CompletionEventID *string `json:"completionEventId"`

}

if parseErr := json.Unmarshal(rpcError.Data, &data); parseErr != nil {

return nil, true, errors.Join(err, parseErr)

}

if deliveryErr := drain(data.CompletionEventID); deliveryErr != nil {

return nil, true, errors.Join(err, deliveryErr)

}

}

}

if cause := context.Cause(requestCtx); cause != nil && errors.Is(err, requestCtx.Err()) {

return nil, true, cause

}

return nil, true, err

}

var completed rpc.SendAndWaitResult

if err := json.Unmarshal(result, &completed); err != nil {

return nil, true, fmt.Errorf("decode message completion: %w", err)

}

if err := drain(completed.CompletionEventID); err != nil {

return nil, true, err

}

if options.ResponseSchema != nil {

if completed.Aborted {

return nil, true, errors.New("session aborted before structured output completed")

}

if completed.Response == nil {

return nil, true, errors.New("run completed without a structured assistant response")

}

data, ok := completed.Response.Data.(*AssistantMessageData)

if !ok || len(data.ToolRequests) > 0 || strings.TrimSpace(data.Content) == "" {

return nil, true, errors.New("run completed without a structured assistant response")

}

}

if completed.Aborted {

return nil, true, errors.New("session aborted before response completed")

}

if completed.Response != nil {

if _, ok := completed.Response.Data.(*AssistantMessageData); !ok ||

(completed.Response.AgentID != nil && *completed.Response.AgentID != "") {

return nil, true, errors.New("the runtime returned a non-assistant response")

}

}

return completed.Response, true, nil

}

func (s *Session) sendAndWaitLegacy(ctx context.Context, options MessageOptions) (*SessionEvent, error) {

type outcome struct {

event *SessionEvent

err error

}

terminalCh := make(chan outcome, 1)

dispatchedCh := make(chan struct{}, 1)

var lastAssistantMessage *SessionEvent

completed := false

var waitMu sync.Mutex

var messageID string

var requestStarted bool

var admitted bool

var observed, dispatched, completionSequence uint64

type recordedEvent struct {

event SessionEvent

sequence uint64

}

var pending []recordedEvent

finish := func(result outcome) {

if !completed {

completed = true

terminalCh <- result

}

}

process := func(event SessionEvent) {

if completed || (event.AgentID != nil && *event.AgentID != "") {

return

}

switch d := event.Data.(type) {

case *UserMessageData:

if messageID != "" && d.MessageID != nil && *d.MessageID == messageID {

requestStarted = true

}

case *SessionWarningData:

if messageID != "" && d.WarningType == "policy_blocked" &&

d.OriginatingMessageID != nil && *d.OriginatingMessageID == messageID {

requestStarted = true

}

case *AssistantMessageData:

if messageID != "" {

if d.OriginatingMessageID != nil && *d.OriginatingMessageID != messageID {

return

}

if d.OriginatingMessageID == nil && !requestStarted {

return

}

}

if messageID != "" && d.OriginatingMessageID != nil && *d.OriginatingMessageID == messageID {

requestStarted = true

}

eventCopy := event

lastAssistantMessage = &eventCopy

case *SessionAutopilotStoppedData:

if messageID != "" && d.RequestID != nil && *d.RequestID == messageID {

switch d.Reason {

case "error":

finish(outcome{err: errors.New("autopilot request failed")})

case "aborted":

finish(outcome{err: errors.New("session aborted before SendAndWait completed")})

default:

finish(outcome{event: lastAssistantMessage})

}

}

case *SessionIdleData:

if d.Mode != nil && *d.Mode == SessionModeAutopilot {

break

}

if !requestStarted {

break

}

if d.Aborted != nil && *d.Aborted {

finish(outcome{err: errors.New("session aborted before SendAndWait completed")})

break

}

finish(outcome{event: lastAssistantMessage})

case *SessionErrorData:

if d.OriginatingMessageID == nil || *d.OriginatingMessageID == "" || *d.OriginatingMessageID == messageID {

finish(outcome{err: fmt.Errorf("session error: %s", d.Message)})

}

}

}

isWaitEvent := func(event SessionEvent) bool {

if event.AgentID != nil && *event.AgentID != "" {

return false

}

switch event.Data.(type) {

case *UserMessageData, *SessionWarningData, *AssistantMessageData,

*SessionAutopilotStoppedData, *SessionIdleData, *SessionErrorData:

return true

default:

return false

}

}

signalDispatched := func() {

if completed && dispatched >= completionSequence {

select {

case dispatchedCh <- struct{}{}:

default:

}

}

}

// Observe the result before user callbacks can retire the current dispatch.

unobserve := s.on(func(event SessionEvent) {

waitMu.Lock()

defer waitMu.Unlock()

observed++

if !isWaitEvent(event) {

return

}

if !admitted {

pending = append(pending, recordedEvent{event, observed})

} else if !completed {

process(event)

if completed {

completionSequence = observed

}

}

}, true)

defer unobserve()

unsubscribe := s.On(func(event SessionEvent) {

waitMu.Lock()

defer waitMu.Unlock()

dispatched = observed

signalDispatched()

})

defer unsubscribe()

id, err := s.Send(ctx, options)

if err != nil {

return nil, err

}

waitMu.Lock()

messageID, admitted = id, true

for _, item := range pending {

process(item.event)

if completed {

completionSequence = item.sequence

break

}

}

pending = nil

signalDispatched()

waitMu.Unlock()

select {

case <-dispatchedCh:

case <-s.eventDone:

s.handlerMutex.RLock()

graceful := s.gracefulRetirement

s.handlerMutex.RUnlock()

if !graceful {

select {

case <-dispatchedCh:

default:

return nil, errors.New("session closed before response completed")

}

}

case <-ctx.Done():

return nil, fmt.Errorf("waiting for response completion: %w", ctx.Err())

}

select {

case result := <-terminalCh:

return result.event, result.err

default:

return nil, errors.New("session closed before response completed")

}

}

// SendPromptAndWait is a convenience wrapper for [Session.SendAndWait] that

// takes a plain prompt string instead of a [MessageOptions] struct. Equivalent

// to:

//

// session.SendAndWait(ctx, copilot.MessageOptions{Prompt: prompt})

func (s *Session) SendPromptAndWait(ctx context.Context, prompt string) (*SessionEvent, error) {

return s.SendAndWait(ctx, MessageOptions{Prompt: prompt})

}

// On subscribes to events from this session.

//

// Events include assistant messages, tool executions, errors, and session state

// changes. Multiple handlers can be registered and will all receive events.

// Handlers are called synchronously in the order they were registered.

//

// The returned function can be called to unsubscribe the handler. It is safe

// to call the unsubscribe function multiple times.

// Unsubscribing also removes an entry from dispatch snapshots if it has not started.

// Registration after the session closes is logged and returns a no-op unsubscribe.

//

// Example:

//

// unsubscribe := session.On(func(event copilot.SessionEvent) {

// switch d := event.Data.(type) {

// case *copilot.AssistantMessageData:

// fmt.Println("Assistant:", d.Content)

// case *copilot.SessionErrorData:

// fmt.Println("Error:", d.Message)

// }

// })

//

// // Later, to stop receiving events:

// unsubscribe()

func (s *Session) On(handler SessionEventHandler) func() {

return s.on(handler, false)

}

func (s *Session) on(handler SessionEventHandler, prepend bool) func() {

s.handlerMutex.Lock()

defer s.handlerMutex.Unlock()

select {

case <-s.eventDone:

log.Printf("Cannot register an event handler on closed session %s", s.SessionID)

return func() {}

default:

}

id := s.nextHandlerID

s.nextHandlerID++

entry := &sessionHandler{id: id, fn: handler}

if prepend {

s.handlers = append([]*sessionHandler{entry}, s.handlers...)

} else {

s.handlers = append(s.handlers, entry)

}

// Return unsubscribe function

return func() {

s.handlerMutex.Lock()

defer s.handlerMutex.Unlock()

for i, h := range s.handlers {

if h.id == id {

h.fn = nil

s.handlers = append(s.handlers[:i], s.handlers[i+1:]...)

break

}

}

}

}

// registerTools registers tool handlers for this session.

//

// Tools with handlers allow the assistant to execute custom functions automatically.

// Declaration-only tools are surfaced as events and left pending for the consumer.

//

// This method is internal and typically called when creating a session with tools.

func (s *Session) registerTools(tools []Tool) {

if err := s.replaceToolHandlers(buildToolHandlerMap(tools)); err != nil {

log.Printf("Cannot register tools on session %s: %v", s.SessionID, err)

}

}

func buildToolHandlerMap(tools []Tool) map[string]ToolHandler {

handlers := make(map[string]ToolHandler)

for _, tool := range tools {

if tool.Name == "" || tool.Handler == nil {

continue

}

handlers[tool.Name] = tool.Handler

}

return handlers

}

func (s *Session) replaceToolHandlers(handlers map[string]ToolHandler) error {

s.toolHandlersM.Lock()

defer s.toolHandlersM.Unlock()

select {

case <-s.eventDone:

return errors.New("session closed before tool handlers could be installed")

default:

}

s.toolHandlers = handlers

return nil

}

// getToolHandler retrieves a registered tool handler by name.

// Returns the handler and true if found, or nil and false if not registered.

func (s *Session) getToolHandler(name string) (ToolHandler, bool) {

s.toolHandlersM.RLock()

handler, ok := s.toolHandlers[name]

s.toolHandlersM.RUnlock()

return handler, ok

}

type permissionHandlerFunc func(

request PermissionRequest,

invocation PermissionInvocation,

promptRequest rpc.PermissionPromptRequest,

) (rpc.PermissionDecision, error)

// registerPermissionHandler registers a permission handler for this session.

//

// When the assistant needs permission to perform certain actions (e.g., file