6cd1733 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

fix(mcp): SSE MCPs not working (#1157)

6cd1733Carlos Alexandro Becker1 year ago

internal/llm/agent/mcp-tools.go+10−7

compare with main

unifiedsplit

@@ -341,29 +341,32 @@ func createAndInitializeClient(ctx context.Context, name string, m config.MCPCon

341 return nil, err

342 }

343

344+ // XXX: ideally we should be able to use context.WithTimeout here, but,

345+ // the SSE MCP client will start failing once that context is canceled.

344 timeout := mcpTimeout(m)

346 timeout := mcpTimeout(m)

345− initCtx, cancel := context.WithTimeout(ctx, timeout)

347+ mcpCtx, cancel := context.WithCancel(ctx)

346− defer cancel()

348+ cancelTimer := time.AfterFunc(timeout, cancel)

347−

348− if err := c.Start(initCtx); err != nil {

349+ if err := c.Start(mcpCtx); err != nil {

349 updateMCPState(name, MCPStateError, maybeTimeoutErr(err, timeout), nil, 0)

350 updateMCPState(name, MCPStateError, maybeTimeoutErr(err, timeout), nil, 0)

350 slog.Error("error starting mcp client", "error", err, "name", name)

351 slog.Error("error starting mcp client", "error", err, "name", name)

351 _ = c.Close()

352 _ = c.Close()

353+ cancel()

352 return nil, err

354 return nil, err

353 }

355 }

354− if _, err := c.Initialize(initCtx, mcpInitRequest); err != nil {

356+ if _, err := c.Initialize(mcpCtx, mcpInitRequest); err != nil {

355 updateMCPState(name, MCPStateError, maybeTimeoutErr(err, timeout), nil, 0)

357 updateMCPState(name, MCPStateError, maybeTimeoutErr(err, timeout), nil, 0)

356 slog.Error("error initializing mcp client", "error", err, "name", name)

358 slog.Error("error initializing mcp client", "error", err, "name", name)

357 _ = c.Close()

359 _ = c.Close()

360+ cancel()

358 return nil, err

361 return nil, err

359 }

362 }

360−

363+ cancelTimer.Stop()

361 slog.Info("Initialized mcp client", "name", name)

364 slog.Info("Initialized mcp client", "name", name)

362 return c, nil

365 return c, nil

363}

366}

364

367

365func maybeTimeoutErr(err error, timeout time.Duration) error {

368func maybeTimeoutErr(err error, timeout time.Duration) error {

366− if errors.Is(err, context.DeadlineExceeded) {

369+ if errors.Is(err, context.Canceled) {

367 return fmt.Errorf("timed out after %s", timeout)

370 return fmt.Errorf("timed out after %s", timeout)

368 }

371 }

369 return err

372 return err