diff --git a/docsgpt/agents/tool_executor.py b/docsgpt/agents/tool_executor.py index 12b36c5b..f4da8263 100644 --- a/docsgpt/agents/tool_executor.py +++ b/docsgpt/agents/tool_executor.py @@ -812,15 +812,24 @@ class ToolExecutor: @staticmethod def _connection_payload(resolved) -> Dict: - """What the chat's Connect card needs; never an account or a secret.""" - return { + """What the chat's Connect card needs; never an account or a secret. + + ``connection_id`` is only the caller's own connection, which the card + reconnects in place; an owner's account (``owner_account``) is not + the caller's to reconnect. + """ + payload = { "connector_key": resolved.connector_key, "connector_name": resolved.connector_name, "status": ( "missing" if resolved.row is None else (resolved.row.get("status") or "reconnect_needed") ), + "owner_account": bool(resolved.delegated), } + if resolved.row is not None and not resolved.delegated: + payload["connection_id"] = str(resolved.row["id"]) + return payload def check_pause(self, tools_dict: Dict, call, llm_class_name: str) -> Optional[Dict]: """Return a pending-action dict (approval / client / headless_denied) or None. diff --git a/docsgpt/connectors/service.py b/docsgpt/connectors/service.py index 7d223b63..b801f1c3 100644 --- a/docsgpt/connectors/service.py +++ b/docsgpt/connectors/service.py @@ -464,6 +464,7 @@ def mark_reconnect_needed(connection_id: str, error: str) -> None: if row is None: return already = normalize_status(row) == STATUS_RECONNECT + counts = repo.resource_counts([connection_id]).get(str(connection_id), {}) repo.update(connection_id, {"status": STATUS_RECONNECT, "last_error": error[:500]}) conn.execute( text( @@ -480,6 +481,8 @@ def mark_reconnect_needed(connection_id: str, error: str) -> None: "connection_id": connection_id, "connector_key": catalog.connector_key_for_row(row), "name": serialize_connection(row)["name"], + "source_count": counts.get("sources", 0), + "tool_count": counts.get("tools", 0), }, scope={"kind": "connection", "id": connection_id}, ) diff --git a/frontend/src/conversation/ConnectToolCallBar.test.tsx b/frontend/src/conversation/ConnectToolCallBar.test.tsx index 2175639a..e5043680 100644 --- a/frontend/src/conversation/ConnectToolCallBar.test.tsx +++ b/frontend/src/conversation/ConnectToolCallBar.test.tsx @@ -20,6 +20,11 @@ vi.mock('../api/services/connectorsService', () => ({ }, })); +const launch = vi.fn(); +vi.mock('../connectors/useConnectorLauncher', () => ({ + default: () => ({ launch, modals: null }), +})); + import connectorsReducer from '../connectors/connectorsSlice'; import ConnectToolCallBar from './ConnectToolCallBar'; import type { ToolCallsType } from './types'; @@ -54,13 +59,31 @@ describe('ConnectToolCallBar', () => { container.remove(); }); - const render = async (toolCall: ToolCallsType, onToolAction = vi.fn()) => { + const render = async ( + toolCall: ToolCallsType, + onToolAction = vi.fn(), + connectors?: Record, + ) => { const store = configureStore({ reducer: { connectors: connectorsReducer, preference: (state = { token: null }) => state, }, - }); + preloadedState: connectors + ? { + connectors: { + enabled: true, + loading: false, + loaded: true, + failed: false, + catalog: [], + connections: [], + ...connectors, + }, + preference: { token: null }, + } + : undefined, + } as Parameters[0]); await act(async () => { root.render( @@ -103,6 +126,47 @@ describe('ConnectToolCallBar', () => { expect(onToolAction).toHaveBeenCalledWith('call-1', 'denied'); }); + const TELEGRAM = { + key: 'telegram', + name: 'Telegram', + icon: 'tool_telegram', + auth_kind: 'api_key', + }; + const ownCall = (status: string): ToolCallsType => ({ + ...call(status), + connection_required: { + connector_key: 'telegram', + connector_name: 'Telegram', + status, + connection_id: 'conn-1', + owner_account: false, + }, + }); + + it('reconnects your own account right here', async () => { + launch.mockClear(); + await render(ownCall('reconnect_needed'), vi.fn(), { + catalog: [TELEGRAM], + connections: [{ id: 'conn-1', status: 'reconnect_needed' }], + }); + await act(async () => + button('conversation.toolApproval.connect:Telegram').click(), + ); + expect(launch).toHaveBeenCalledWith(TELEGRAM, { + mode: 'reconnect', + connectionId: 'conn-1', + }); + expect(document.body.textContent).not.toContain('DRAWER'); + }); + + it('offers Continue once the connection works again', async () => { + await render(ownCall('reconnect_needed'), vi.fn(), { + catalog: [TELEGRAM], + connections: [{ id: 'conn-1', status: 'connected' }], + }); + expect(button('conversation.toolApproval.continue')).toBeDefined(); + }); + it('an account that needs signing in again is healed from its drawer', async () => { await render(call('reconnect_needed')); await act(async () => diff --git a/frontend/src/conversation/ConnectToolCallBar.tsx b/frontend/src/conversation/ConnectToolCallBar.tsx index 462fb372..eb12bd0e 100644 --- a/frontend/src/conversation/ConnectToolCallBar.tsx +++ b/frontend/src/conversation/ConnectToolCallBar.tsx @@ -8,6 +8,7 @@ import { Button } from '../components/ui/button'; import ConnectorIcon from '../connectors/ConnectorIcon'; import { loadConnectors, + selectConnections, selectConnectorCatalog, selectConnectorsLoaded, } from '../connectors/connectorsSlice'; @@ -34,12 +35,18 @@ export default function ConnectToolCallBar({ const token = useSelector(selectToken); const catalog = useSelector(selectConnectorCatalog); const loaded = useSelector(selectConnectorsLoaded); + const connections = useSelector(selectConnections); const [connected, setConnected] = useState(false); const { launch, modals } = useConnectorLauncher({ onConnected: () => setConnected(true), }); const required = toolCall.connection_required; const connector = catalog.find((c) => c.key === required?.connector_key); + const ownConnection = required?.connection_id + ? connections.find((c) => c.id === required.connection_id) + : undefined; + // Connected here, or healthy again when the user comes back to the chat. + const ready = connected || ownConnection?.status === 'connected'; const name = required?.connector_name || connector?.name || @@ -50,13 +57,22 @@ export default function ConnectToolCallBar({ }, [loaded, dispatch, token]); const connect = () => { - // A member who has no account for the service connects one. An owner's - // existing account that needs signing in again is healed from its drawer, - // which reconnects that same connection for every tool and source. + // No account yet: connect one. The caller's own account that needs + // signing in again is reconnected right here (every tool and source on + // it heals). MCP servers and anyone else's account go to the drawer. if (connector && required?.status === 'missing') { launch(connector); return; } + const inPlace = + connector?.auth_kind === 'oauth' || connector?.auth_kind === 'api_key'; + if (connector && required?.connection_id && inPlace) { + launch(connector, { + mode: 'reconnect', + connectionId: required.connection_id, + }); + return; + } navigate( `/settings/connectors${connector ? `?connector=${encodeURIComponent(connector.key)}` : ''}`, ); @@ -78,7 +94,7 @@ export default function ConnectToolCallBar({
- {connected ? ( + {ready ? (